|
| 1 | +// Copyright 2021 - 2022 Matrix Origin |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +package hakeeper |
| 16 | + |
| 17 | +import ( |
| 18 | + "encoding/binary" |
| 19 | + "encoding/gob" |
| 20 | + "io" |
| 21 | + "time" |
| 22 | + |
| 23 | + "github.com/lni/dragonboat/v4/logger" |
| 24 | + sm "github.com/lni/dragonboat/v4/statemachine" |
| 25 | + |
| 26 | + "github.com/matrixorigin/matrixone/pkg/common/moerr" |
| 27 | + "github.com/matrixorigin/matrixone/pkg/pb/logservice" |
| 28 | +) |
| 29 | + |
| 30 | +var ( |
| 31 | + plog = logger.GetLogger("hakeeper") |
| 32 | +) |
| 33 | + |
| 34 | +var ( |
| 35 | + binaryEnc = binary.BigEndian |
| 36 | +) |
| 37 | + |
| 38 | +const ( |
| 39 | + // TickDuration defines the frequency of ticks. |
| 40 | + TickDuration = time.Second |
| 41 | + // DefaultHAKeeperShardID is the shard ID assigned to the special HAKeeper |
| 42 | + // shard. |
| 43 | + DefaultHAKeeperShardID uint64 = 0 |
| 44 | + |
| 45 | + headerSize = 2 |
| 46 | +) |
| 47 | + |
| 48 | +const ( |
| 49 | + createLogShardTag uint16 = iota + 0xAE01 |
| 50 | + tickTag |
| 51 | + dnHeartbeatTag |
| 52 | + logHeartbeatTag |
| 53 | +) |
| 54 | + |
| 55 | +type logShardIDQuery struct { |
| 56 | + name string |
| 57 | +} |
| 58 | + |
| 59 | +type logShardIDQueryResult struct { |
| 60 | + id uint64 |
| 61 | + found bool |
| 62 | +} |
| 63 | + |
| 64 | +type stateMachine struct { |
| 65 | + replicaID uint64 |
| 66 | + |
| 67 | + Tick uint64 |
| 68 | + NextID uint64 |
| 69 | + |
| 70 | + LogShards map[string]uint64 |
| 71 | + DNState DNState |
| 72 | + LogState LogState |
| 73 | +} |
| 74 | + |
| 75 | +func parseCmdTag(cmd []byte) uint16 { |
| 76 | + return binaryEnc.Uint16(cmd) |
| 77 | +} |
| 78 | + |
| 79 | +func getCreateLogShardCmd(name string) []byte { |
| 80 | + return getLogShardCmd(name, createLogShardTag) |
| 81 | +} |
| 82 | + |
| 83 | +func getLogShardCmd(name string, tag uint16) []byte { |
| 84 | + cmd := make([]byte, headerSize+len(name)) |
| 85 | + binaryEnc.PutUint16(cmd, tag) |
| 86 | + copy(cmd[headerSize:], []byte(name)) |
| 87 | + return cmd |
| 88 | +} |
| 89 | + |
| 90 | +func isCreateLogShardCmd(cmd []byte) (string, bool) { |
| 91 | + return isLogShardCmd(cmd, createLogShardTag) |
| 92 | +} |
| 93 | + |
| 94 | +func isDNHeartbeatCmd(cmd []byte) bool { |
| 95 | + return isHeartbeatCmd(cmd, dnHeartbeatTag) |
| 96 | +} |
| 97 | + |
| 98 | +func isLogHeartbeatCmd(cmd []byte) bool { |
| 99 | + return isHeartbeatCmd(cmd, logHeartbeatTag) |
| 100 | +} |
| 101 | + |
| 102 | +func isHeartbeatCmd(cmd []byte, tag uint16) bool { |
| 103 | + if len(cmd) <= headerSize { |
| 104 | + return false |
| 105 | + } |
| 106 | + return parseCmdTag(cmd) == tag |
| 107 | +} |
| 108 | + |
| 109 | +func parseHeartbeatCmd(cmd []byte) []byte { |
| 110 | + return cmd[headerSize:] |
| 111 | +} |
| 112 | + |
| 113 | +func isLogShardCmd(cmd []byte, tag uint16) (string, bool) { |
| 114 | + if len(cmd) <= headerSize { |
| 115 | + return "", false |
| 116 | + } |
| 117 | + if parseCmdTag(cmd) == tag { |
| 118 | + return string(cmd[headerSize:]), true |
| 119 | + } |
| 120 | + return "", false |
| 121 | +} |
| 122 | + |
| 123 | +func isTickCmd(cmd []byte) bool { |
| 124 | + return len(cmd) == headerSize && binaryEnc.Uint16(cmd) == tickTag |
| 125 | +} |
| 126 | + |
| 127 | +func GetTickCmd() []byte { |
| 128 | + cmd := make([]byte, headerSize) |
| 129 | + binaryEnc.PutUint16(cmd, tickTag) |
| 130 | + return cmd |
| 131 | +} |
| 132 | + |
| 133 | +func GetLogStoreHeartbeatCmd(data []byte) []byte { |
| 134 | + return getHeartbeatCmd(data, logHeartbeatTag) |
| 135 | +} |
| 136 | + |
| 137 | +func GetDNStoreHeartbeatCmd(data []byte) []byte { |
| 138 | + return getHeartbeatCmd(data, dnHeartbeatTag) |
| 139 | +} |
| 140 | + |
| 141 | +func getHeartbeatCmd(data []byte, tag uint16) []byte { |
| 142 | + cmd := make([]byte, headerSize+len(data)) |
| 143 | + binaryEnc.PutUint16(cmd, tag) |
| 144 | + copy(cmd[headerSize:], data) |
| 145 | + return cmd |
| 146 | +} |
| 147 | + |
| 148 | +func NewStateMachine(shardID uint64, replicaID uint64) sm.IStateMachine { |
| 149 | + if shardID != DefaultHAKeeperShardID { |
| 150 | + panic(moerr.NewError(moerr.INVALID_INPUT, "invalid HAKeeper shard ID")) |
| 151 | + } |
| 152 | + return &stateMachine{ |
| 153 | + replicaID: replicaID, |
| 154 | + LogShards: make(map[string]uint64), |
| 155 | + DNState: NewDNState(), |
| 156 | + LogState: NewLogState(), |
| 157 | + } |
| 158 | +} |
| 159 | + |
| 160 | +func (s *stateMachine) Close() error { |
| 161 | + return nil |
| 162 | +} |
| 163 | + |
| 164 | +func (s *stateMachine) assignID() uint64 { |
| 165 | + s.NextID++ |
| 166 | + return s.NextID |
| 167 | +} |
| 168 | + |
| 169 | +func (s *stateMachine) handleCreateLogShardCmd(cmd []byte) (sm.Result, error) { |
| 170 | + name, ok := isCreateLogShardCmd(cmd) |
| 171 | + if !ok { |
| 172 | + panic(moerr.NewError(moerr.INVALID_INPUT, "not create log shard cmd")) |
| 173 | + } |
| 174 | + if shardID, ok := s.LogShards[name]; ok { |
| 175 | + data := make([]byte, 8) |
| 176 | + binaryEnc.PutUint64(data, shardID) |
| 177 | + return sm.Result{Value: 0, Data: data}, nil |
| 178 | + } |
| 179 | + s.LogShards[name] = s.assignID() |
| 180 | + return sm.Result{Value: s.NextID}, nil |
| 181 | +} |
| 182 | + |
| 183 | +func (s *stateMachine) handleDNHeartbeat(cmd []byte) (sm.Result, error) { |
| 184 | + data := parseHeartbeatCmd(cmd) |
| 185 | + var hb logservice.DNStoreHeartbeat |
| 186 | + if err := hb.Unmarshal(data); err != nil { |
| 187 | + panic(err) |
| 188 | + } |
| 189 | + s.DNState.Update(hb, s.Tick) |
| 190 | + return sm.Result{}, nil |
| 191 | +} |
| 192 | + |
| 193 | +func (s *stateMachine) handleLogHeartbeat(cmd []byte) (sm.Result, error) { |
| 194 | + data := parseHeartbeatCmd(cmd) |
| 195 | + var hb logservice.LogStoreHeartbeat |
| 196 | + if err := hb.Unmarshal(data); err != nil { |
| 197 | + panic(err) |
| 198 | + } |
| 199 | + s.LogState.Update(hb, s.Tick) |
| 200 | + return sm.Result{}, nil |
| 201 | +} |
| 202 | + |
| 203 | +func (s *stateMachine) handleTick(cmd []byte) (sm.Result, error) { |
| 204 | + s.Tick++ |
| 205 | + return sm.Result{}, nil |
| 206 | +} |
| 207 | + |
| 208 | +func (s *stateMachine) Update(e sm.Entry) (sm.Result, error) { |
| 209 | + cmd := e.Cmd |
| 210 | + if _, ok := isCreateLogShardCmd(cmd); ok { |
| 211 | + return s.handleCreateLogShardCmd(cmd) |
| 212 | + } else if isDNHeartbeatCmd(cmd) { |
| 213 | + return s.handleDNHeartbeat(cmd) |
| 214 | + } else if isLogHeartbeatCmd(cmd) { |
| 215 | + return s.handleLogHeartbeat(cmd) |
| 216 | + } else if isTickCmd(cmd) { |
| 217 | + return s.handleTick(cmd) |
| 218 | + } |
| 219 | + panic(moerr.NewError(moerr.INVALID_INPUT, "unexpected haKeeper cmd")) |
| 220 | +} |
| 221 | + |
| 222 | +func (s *stateMachine) Lookup(query interface{}) (interface{}, error) { |
| 223 | + if q, ok := query.(*logShardIDQuery); ok { |
| 224 | + id, ok := s.LogShards[q.name] |
| 225 | + if ok { |
| 226 | + return &logShardIDQueryResult{found: true, id: id}, nil |
| 227 | + } |
| 228 | + return &logShardIDQueryResult{found: false}, nil |
| 229 | + } |
| 230 | + panic("unknown query type") |
| 231 | +} |
| 232 | + |
| 233 | +func (s *stateMachine) SaveSnapshot(w io.Writer, |
| 234 | + _ sm.ISnapshotFileCollection, _ <-chan struct{}) error { |
| 235 | + enc := gob.NewEncoder(w) |
| 236 | + return enc.Encode(s) |
| 237 | +} |
| 238 | + |
| 239 | +func (s *stateMachine) RecoverFromSnapshot(r io.Reader, |
| 240 | + _ []sm.SnapshotFile, _ <-chan struct{}) error { |
| 241 | + dec := gob.NewDecoder(r) |
| 242 | + return dec.Decode(s) |
| 243 | +} |
0 commit comments