llamerada-jp / colonio / 18451467117

func NewKVS(conf *Config) *KVS {

        k := &KVS{

                logger:     conf.Logger,

                handler:    conf.Handler,

                store:      conf.Store,

                transferer: conf.Transferer,

                sectors:    make(map[config.KvsSectorKey]*sector),

                members:    make(map[uint64]*memberEntry),

        }

        transferer.SetRequestHandler[proto.PacketContent_KvsOperation](k.transferer, k.recvOperation)

        transferer.SetRequestHandler[proto.PacketContent_RaftConfig](k.transferer, k.recvConfig)

        transferer.SetRequestHandler[proto.PacketContent_RaftConfigResponse](k.transferer, k.recvConfigResponse)

        transferer.SetRequestHandler[proto.PacketContent_RaftMessage](k.transferer, k.recvMessage)

        return k

}

func (k *KVS) Start(ctx context.Context, localNodeID *shared.NodeID) {

        k.localNodeID = localNodeID

        k.ctx = ctx

        go func() {

                ticker := time.NewTicker(1 * time.Second)

                defer ticker.Stop()

                for {

                        select {

                        case <-ctx.Done():

                                return

                        case <-ticker.C:

                                k.subRoutine()

func (k *KVS) Get(key string) chan *config.KvsGetResult {

        c := make(chan *config.KvsGetResult, 1)

        k.sendOperation(proto.KvsOperation_COMMAND_GET, key, nil, func(res *proto.KvsOperationResponse, err error) {

                defer close(c)

                if err != nil {

                        c <- &config.KvsGetResult{

                                Data: nil,

                                Err:  err,

                        }

                        return

                }

                switch res.Error {

                case proto.KvsOperationResponse_ERROR_NONE:

                        // ok

                        c <- &config.KvsGetResult{

                                Data: res.Value,

                                Err:  nil,

                        }

                case proto.KvsOperationResponse_ERROR_NOT_FOUND:

                        c <- &config.KvsGetResult{

                                Data: nil,

                                Err:  fmt.Errorf("key not found: %s", key),

                        }

                default:

                        c <- &config.KvsGetResult{

                                Data: nil,

                                Err:  fmt.Errorf("unknown error: %d", res.Error),

                        }

        return c

func (k *KVS) Set(key string, value []byte) chan error {

        c := make(chan error, 1)

        k.sendOperation(proto.KvsOperation_COMMAND_SET, key, value, func(res *proto.KvsOperationResponse, err error) {

                defer close(c)

                if err != nil {

                        c <- err

                        return

                }

                if res.Error != proto.KvsOperationResponse_ERROR_NONE {

                        c <- fmt.Errorf("kvs set error: %d", res.Error)

                        return

                }

                c <- nil

        return c

func (k *KVS) Patch(key string, value []byte) chan error {

        c := make(chan error, 1)

        k.sendOperation(proto.KvsOperation_COMMAND_PATCH, key, value, func(res *proto.KvsOperationResponse, err error) {

                defer close(c)

                if err != nil {

                        c <- err

                        return

                }

                if res.Error != proto.KvsOperationResponse_ERROR_NONE {

                        c <- fmt.Errorf("kvs patch error: %d", res.Error)

                        return

                }

                c <- nil

        return c

func (k *KVS) Delete(key string) chan error {

        c := make(chan error, 1)

        k.sendOperation(proto.KvsOperation_COMMAND_DELETE, key, nil, func(res *proto.KvsOperationResponse, err error) {

                defer close(c)

                if err != nil {

                        c <- err

                        return

                }

                if res.Error != proto.KvsOperationResponse_ERROR_NONE {

                        c <- fmt.Errorf("kvs delete error: %d", res.Error)

                        return

                }

                c <- nil

        return c

func (k *KVS) subRoutine() {

        isStable, nextNodeIDs := k.handler.KvsGetStability()

        if !isStable {

                return

        }

        k.mtx.Lock()

        defer k.mtx.Unlock()

        if k.coordinating == nil {

                k.createCluster(nextNodeIDs)

                if k.coordinating == nil {

                        return

                }

        toAppend, toRemove := k.getNodeDiff(nextNodeIDs)

        for _, nodeID := range toAppend {

                k.lastSequence++

                k.members[k.lastSequence] = &memberEntry{

                        nodeID: nodeID,

                        state:  memberStateAppendingRaft,

                }

        }

        for seq := range toRemove {

                k.members[seq].state = memberStateRemoving

        }

        if err := k.applyConfigRaft(); err != nil {

                k.logger.Warn("Failed to apply Raft config", "error", err)

        }

        k.sendSettingMessage()

func (k *KVS) getNodeDiff(nextNodeIDs []*shared.NodeID) ([]*shared.NodeID, map[uint64]*memberEntry) {

        nextNodeIDMap := make(map[shared.NodeID]struct{})

        for _, nodeID := range nextNodeIDs {

                nextNodeIDMap[*nodeID] = struct{}{}

        }

        memberMap := make(map[shared.NodeID]uint64)

        for seq, entry := range k.members {

                if entry.state == memberStateRemoving {

                        continue

                memberMap[*entry.nodeID] = seq

        toAppend := make([]*shared.NodeID, 0)

        toRemove := make(map[uint64]*memberEntry)

        for nodeID := range nextNodeIDMap {

                if _, ok := memberMap[nodeID]; !ok {

                        toAppend = append(toAppend, &nodeID)

                }

        for seq, entry := range k.members {

                if entry.state == memberStateRemoving {

                        continue

                if _, ok := nextNodeIDMap[*entry.nodeID]; !ok {

                        toRemove[seq] = entry

                }

        return toAppend, toRemove

func (k *KVS) createCluster(nextNodeIDs []*shared.NodeID) {

        clusterID, err := uuid.NewV7()

        if err != nil {

                k.logger.Error("Failed to create new cluster ID", "error", err)

                return

        }

        k.lastSequence = coordinatingNodeSequence

        member := make(map[uint64]*shared.NodeID)

        member[coordinatingNodeSequence] = k.localNodeID

        for _, nodeID := range nextNodeIDs {

                k.lastSequence++

                member[k.lastSequence] = nodeID

                k.members[k.lastSequence] = &memberEntry{

                        nodeID: nodeID,

                        state:  memberStateCreating,

                }

        }

        sectorKey := &config.KvsSectorKey{

                SectorID: clusterID,

                Sequence: coordinatingNodeSequence,

        }

        k.coordinating = sectorKey

        k.allocateCluster(sectorKey, k.localNodeID, false, member)

func (k *KVS) allocateCluster(sectorKey *config.KvsSectorKey, head *shared.NodeID, append bool, members map[uint64]*shared.NodeID) {

        s := &sector{}

        k.sectors[*sectorKey] = s

        raftNode := newRaftNode(&raftNodeConfig{

                logger:    k.logger,

                manager:   k,

                store:     s,

                sectorKey: sectorKey,

                join:      append,

                member:    members,

        })

        store := newStore(&storeConfig{

                sectorKey: sectorKey,

                handler:   s,

                head:      head,

        })

        s.raft = raftNode

        s.store = store

        raftNode.start(k.ctx)

}

func (k *KVS) applyConfigRaft() error {

        raftNode := k.sectors[*k.coordinating].raft

        for sequence, member := range k.members {

                switch member.state {

                case memberStateAppendingRaft:

                        return raftNode.appendNode(sequence, member.nodeID)

                case memberStateRemoving:

                        return raftNode.removeNode(sequence)

        return nil

func (h *kvsOperationHandler) OnResponse(packet *shared.Packet) {

        h.receiver(packet.Content.GetKvsOperationResponse(), nil)

}

func (h *kvsOperationHandler) OnError(code constants.PacketErrorCode, message string) {

        h.receiver(nil, fmt.Errorf("packet error %d: %s", code, message))

}

func (k *KVS) sendOperation(command proto.KvsOperation_Command, key string, value []byte, receiver func(res *proto.KvsOperationResponse, err error)) {

        dst := shared.NewHashedNodeID([]byte(key))

        content := &proto.PacketContent{

                Content: &proto.PacketContent_KvsOperation{

                        KvsOperation: &proto.KvsOperation{

                                Command: command,

                                Key:     key,

                                Value:   value,

                        },

                },

        }

        k.transferer.Request(dst, shared.PacketModeNone, content, &kvsOperationHandler{

                receiver: receiver,

        })

}

func (k *KVS) sendSettingMessage() {

        for sequence, member := range k.members {

                if member.state == memberStateRemoving {

                        continue

                members := make(map[uint64]*proto.NodeID)

                members[coordinatingNodeSequence] = k.localNodeID.Proto()

                for ms, mm := range k.members {

                        members[ms] = mm.nodeID.Proto()

                }

                var command proto.RaftConfig_Command

                if member.state == memberStateCreating {

                        command = proto.RaftConfig_COMMAND_CREATE

                } else {

                        command = proto.RaftConfig_COMMAND_APPEND

                }

                k.transferer.RequestOneWay(

                        member.nodeID,

                        shared.PacketModeExplicit|shared.PacketModeNoRetry,

                        &proto.PacketContent{

                                Content: &proto.PacketContent_RaftConfig{

                                        RaftConfig: &proto.RaftConfig{

                                                SectorId: MustMarshalUUID(k.coordinating.SectorID),

                                                Sequence: sequence,

                                                Command:  command,

                                                Members:  members,

                                        },

                                },

                        },

                )

func (k *KVS) recvOperation(packet *shared.Packet) {

        content := packet.Content.GetKvsOperation()

        command := content.Command

        key := content.Key

        value := content.Value

        store := k.sectors[*k.coordinating].store

        switch command {

        case proto.KvsOperation_COMMAND_GET:

                data, err := store.get(key)

                e := proto.KvsOperationResponse_ERROR_NONE

                if err != nil {

                        if err == config.ErrorKvsStoreKeyNotFound {

                                e = proto.KvsOperationResponse_ERROR_NOT_FOUND

                        } else {

                                e = proto.KvsOperationResponse_ERROR_UNKNOWN

                        }

                k.transferer.Response(packet, &proto.PacketContent{

                        Content: &proto.PacketContent_KvsOperationResponse{

                                KvsOperationResponse: &proto.KvsOperationResponse{

                                        Error: e,

                                        Value: data,

                                },

                        },

                })

        case proto.KvsOperation_COMMAND_SET:

                err := store.set(key, value)

                e := proto.KvsOperationResponse_ERROR_NONE

                if err != nil {

                        e = proto.KvsOperationResponse_ERROR_UNKNOWN

                }

                k.transferer.Response(packet, &proto.PacketContent{

                        Content: &proto.PacketContent_KvsOperationResponse{

                                KvsOperationResponse: &proto.KvsOperationResponse{

                                        Error: e,

                                },

                        },

                })

        case proto.KvsOperation_COMMAND_PATCH:

                err := store.patch(key, value)

                e := proto.KvsOperationResponse_ERROR_NONE

                if err != nil {

                        e = proto.KvsOperationResponse_ERROR_UNKNOWN

                }

                k.transferer.Response(packet, &proto.PacketContent{

                        Content: &proto.PacketContent_KvsOperationResponse{

                                KvsOperationResponse: &proto.KvsOperationResponse{

                                        Error: e,

                                },

                        },

                })

        case proto.KvsOperation_COMMAND_DELETE:

                err := store.delete(key)

                e := proto.KvsOperationResponse_ERROR_NONE

                if err != nil {

                        e = proto.KvsOperationResponse_ERROR_UNKNOWN

                }

                k.transferer.Response(packet, &proto.PacketContent{

                        Content: &proto.PacketContent_KvsOperationResponse{

                                KvsOperationResponse: &proto.KvsOperationResponse{

                                        Error: e,

                                },

                        },

                })

        default:

                k.transferer.Response(packet, &proto.PacketContent{

                        Content: &proto.PacketContent_KvsOperationResponse{

                                KvsOperationResponse: &proto.KvsOperationResponse{

                                        Error: proto.KvsOperationResponse_ERROR_UNKNOWN,

                                },

                        },

                })

func (k *KVS) recvConfig(packet *shared.Packet) {

        content := packet.Content.GetRaftConfig()

        sectorID, err := uuid.ParseBytes(content.SectorId)

        if err != nil {

                k.logger.Warn("Failed to parse promoter NodeID", "error", err)

                return

        }

        sequence := content.Sequence

        command := content.Command

        members := make(map[uint64]*shared.NodeID)

        for seq, nodeID := range content.Members {

                var err error

                members[seq], err = shared.NewNodeIDFromProto(nodeID)

                if err != nil {

                        k.logger.Warn("Failed to create NodeID from proto", "error", err, "nodeID", nodeID)

                        return

                }

        k.mtx.Lock()

        defer k.mtx.Unlock()

        sectorKey := config.KvsSectorKey{

                SectorID: sectorID,

                Sequence: sequence,

        }

        switch command {

        case proto.RaftConfig_COMMAND_CREATE, proto.RaftConfig_COMMAND_APPEND:

                if _, ok := k.sectors[sectorKey]; !ok {

                        append := true

                        if command == proto.RaftConfig_COMMAND_CREATE {

                                append = false

                        }

                        k.allocateCluster(&sectorKey, packet.SrcNodeID, append, members)

        default:

                k.logger.Warn("Unknown RaftConfig command", "command", command)

                return

        response := &proto.PacketContent{

                Content: &proto.PacketContent_RaftConfigResponse{

                        RaftConfigResponse: &proto.RaftConfigResponse{

                                SectorId: MustMarshalUUID(sectorID),

                                Sequence: sequence,

                        },

                },

        }

        k.transferer.RequestOneWay(packet.SrcNodeID, shared.PacketModeExplicit|shared.PacketModeNoRetry, response)

func (k *KVS) recvConfigResponse(packet *shared.Packet) {

        content := packet.Content.GetRaftConfigResponse()

        sectorID, err := uuid.ParseBytes(content.SectorId)

        if err != nil {

                k.logger.Warn("Failed to parse promoter NodeID", "error", err)

                return

        }

        sequence := content.Sequence

        k.mtx.Lock()

        defer k.mtx.Unlock()

        // ignore response if it does not match the current node

        if k.coordinating == nil || k.coordinating.SectorID != sectorID {

                return

        }

        if entry, ok := k.members[sequence]; ok && entry.nodeID.Equal(packet.SrcNodeID) {

                entry.state = memberStateNormal

        }

func (k *KVS) recvMessage(packet *shared.Packet) {

        content := packet.Content.GetRaftMessage()

        sectorID, err := uuid.ParseBytes(content.SectorId)

        if err != nil {

                k.logger.Warn("Failed to parse promoter NodeID", "error", err)

                return

        }

        key := config.KvsSectorKey{

                SectorID: sectorID,

                Sequence: content.Sequence,

        }

        k.mtx.RLock()

        n, ok := k.sectors[key]

        k.mtx.RUnlock()

        if !ok {

                return

        }

        if err := n.raft.processMessage(content); err != nil {

                k.logger.Error("Failed to process Raft message", "error", err)

        }

func (k *KVS) raftNodeError(sectorKey *config.KvsSectorKey, err error) {

        panic(fmt.Sprintf("raftNodeError not implemented: %s", err))

func (k *KVS) raftNodeSendMessage(dstNodeID *shared.NodeID, message *proto.RaftMessage) {

        k.transferer.RequestOneWay(

                dstNodeID,

                shared.PacketModeExplicit|shared.PacketModeNoRetry,

                &proto.PacketContent{

                        Content: &proto.PacketContent_RaftMessage{

                                RaftMessage: message,

                        },

                })

}

func (k *KVS) raftNodeApplyProposal(sectorKey *config.KvsSectorKey, proposal *proto.RaftProposalManagement) {

        _, ok := k.sectors[*sectorKey]

        if !ok {

                return

        }

        panic("raftNodeApplyProposal not implemented")

func (k *KVS) raftNodeAppendNode(sectorKey *config.KvsSectorKey, sequence uint64, nodeID *shared.NodeID) {

        if *sectorKey != *k.coordinating {

                return

        }

        k.mtx.Lock()

        defer k.mtx.Unlock()

        member, ok := k.members[sequence]

        // unknown

        if !ok {

                k.logger.Warn("Unknown node sequence for appending", "sequence", sequence, "nodeID", nodeID)

                k.members[sequence] = &memberEntry{

                        nodeID: nodeID,

                        state:  memberStateRemoving,

                }

                return

        }

        if member.state == memberStateAppendingRaft {

                member.state = memberStateAppendingNode

        } else {

                k.logger.Warn("Unexpected state for appending node", "sequence", sequence, "state", member.state, "nodeID", nodeID)

        }

func (k *KVS) raftNodeRemoveNode(sectorKey *config.KvsSectorKey, sequence uint64) {

        k.mtx.RLock()

        sector, ok := k.sectors[*sectorKey]

        k.mtx.RUnlock()

        if !ok {

                return

        }

        if *sectorKey == *k.coordinating {

                delete(k.members, sequence)

        }

        sector.raft.stop()

        k.mtx.Lock()

        defer k.mtx.Unlock()

        delete(k.sectors, *sectorKey)