JCS-pub/client/internal/cluster/fsm.go

package cluster

import (
	"encoding/binary"
	"fmt"
	"io"

	"github.com/hashicorp/raft"
	"gitlink.org.cn/cloudream/common/utils/http2"
	"gitlink.org.cn/cloudream/common/utils/serder"
)

var ErrFSMNotFound = fmt.Errorf("fsm not found")

type FSM interface {
	// 必须唯一且不能变化
	ID() string
	Apply(cmd []byte) ([]byte, error)
	Snapshot() (FSMSnapshot, error)
	Restore(input io.Reader) error
}

type FSMSnapshot interface {
	Persist(output io.Writer) error
	Release()
}

type SnapshotMeta struct {
	Version int
}

type raftFSM struct {
	fsms map[string]FSM
}

func NewFSM(fsms []FSM) *raftFSM {
	fsmsMp := make(map[string]FSM)
	for _, fsm := range fsms {
		_, ok := fsmsMp[fsm.ID()]
		if ok {
			panic(fmt.Sprintf("duplicate fsm id: %s", fsm.ID()))
		}
		fsmsMp[fsm.ID()] = fsm
	}

	return &raftFSM{fsms: fsmsMp}
}

func (f *raftFSM) Apply(l *raft.Log) interface{} {
	idLen := binary.LittleEndian.Uint32(l.Data[:4])
	fsm, ok := f.fsms[string(l.Data[4:4+idLen])]
	if !ok {
		return ErrFSMNotFound
	}

	cmd := l.Data[4+idLen:]
	res, err := fsm.Apply(cmd)
	return applyResult{
		Result: res,
		Error:  err,
	}
}

func (f *raftFSM) Snapshot() (raft.FSMSnapshot, error) {
	snapshots := make(map[string]FSMSnapshot)
	for id, fsm := range f.fsms {
		snapshot, err := fsm.Snapshot()
		if err != nil {
			for _, snapshot := range snapshots {
				snapshot.Release()
			}
			return nil, err
		}

		snapshots[id] = snapshot
	}

	return &Snapshot{snapshots: snapshots}, nil
}

func (f *raftFSM) Restore(rc io.ReadCloser) error {
	defer rc.Close()

	cr := http2.NewChunkedReader(rc)

	_, metaBytes, err := cr.NextDataPart()
	if err != nil {
		return fmt.Errorf("read meta data part: %v", err)
	}
	meta := SnapshotMeta{}
	if err := serder.JSONToObject(metaBytes, &meta); err != nil {
		return fmt.Errorf("unmarshal meta data: %v", err)
	}
	// 进行类似的检查
	if meta.Version != 1 {
		return fmt.Errorf("unsupported version: %d", meta.Version)
	}

	for {
		id, reader, err := cr.NextPart()
		if err != nil && err != io.EOF {
			return err
		}

		if err == io.EOF {
			// TODO 考虑检查一下是否调用了所有FSM的Restore方法
			break
		}

		fsm, ok := f.fsms[id]
		if !ok {
			// TODO 兼容性
			continue
		}

		err = fsm.Restore(reader)
		if err != nil {
			// TODO 不知道Raft库在发现Restore失败后是否能够及时停止服务
			return fmt.Errorf("restore fsm %s: %v", id, err)
		}
	}

	return nil
}

var _ raft.FSM = (*raftFSM)(nil)

type Snapshot struct {
	snapshots map[string]FSMSnapshot
}

func (s *Snapshot) Persist(sink raft.SnapshotSink) error {
	meta := SnapshotMeta{Version: 1}
	metaBytes, err := serder.ObjectToJSON(meta)
	if err != nil {
		sink.Cancel()
		return fmt.Errorf("marshal meta data: %v", err)
	}

	cw := http2.NewChunkedWriter(sink)
	err = cw.WriteDataPart("meta", metaBytes)
	if err != nil {
		sink.Cancel()
		return fmt.Errorf("write meta data part: %v", err)
	}

	for id, snapshot := range s.snapshots {
		w := cw.BeginPart(id)
		err := snapshot.Persist(w)
		if err != nil {
			sink.Cancel()
			return fmt.Errorf("persist fsm %s: %v", id, err)
		}
	}

	return sink.Close()
}

func (s *Snapshot) Release() {
	for _, snapshot := range s.snapshots {
		snapshot.Release()
	}
}

type applyResult struct {
	Result []byte
	Error  error
}