fusen-render/sm.go

package fusenrender

import (
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"sync"
	"time"

	sm "github.com/lni/dragonboat/v4/statemachine"
)

type SMQueue struct {
	shardID   uint64
	replicaID uint64

	mu     sync.Mutex
	queues map[string]*PriorityQueue[QueueItem]
}

type Command struct {
	Name  string `json:"name"`
	Group string `json:"group"`

	Item *QueueItem `json:"item"`
}

func (cmd *Command) Encode() ([]byte, error) {
	val, err := json.Marshal(cmd)
	if err != nil {
		return nil, err
	}
	return val, nil
}

type QueueItem struct {
	Group    string    `json:"group"`     // 组名
	Priority uint32    `json:"priority"`  // 处理的优先级
	CreateAt time.Time `json:"create_at"` // 创建时间 统一utc
	Data     any       `json:"data"`      // 操作的数据结构
}

func (item *QueueItem) Encode() ([]byte, error) {
	val, err := json.Marshal(item)
	if err != nil {
		return nil, err
	}
	return val, nil
}

func (item *QueueItem) GetKey() []byte {
	return []byte(fmt.Sprintf("%s_%d_%d", item.Group, -item.Priority, item.CreateAt.UTC().Unix()))
}

// // NewSMQueue creates and return a new ExampleStateMachine object.
func NewSMQueue(shardID uint64, replicaID uint64) sm.IStateMachine {

	return &SMQueue{
		shardID:   shardID,
		replicaID: replicaID,

		queues: make(map[string]*PriorityQueue[QueueItem]),
	}
}

// Lookup performs local lookup on the ExampleStateMachine instance. In this example,
// we always return the Count value as a little endian binary encoded byte
// slice.
func (s *SMQueue) Lookup(group interface{}) (item interface{}, err error) {

	return item, nil
}

// Update updates the object using the specified committed raft entry.
func (s *SMQueue) Update(e sm.Entry) (result sm.Result, err error) {

	var cmd Command

	err = json.Unmarshal(e.Cmd, &cmd)
	if err != nil {
		return result, err
	}

	switch cmd.Name {
	case "enqueue":
		d, err := cmd.Item.Encode()
		if err != nil {
			return result, err
		}
		key := cmd.Item.GetKey()

		s.mu.Lock()
		var queue *PriorityQueue[QueueItem]
		var ok bool
		if queue, ok = s.queues[cmd.Group]; !ok {
			queue = NewPriorityQueue[QueueItem]()
			s.queues[cmd.Group] = queue
		}
		defer s.mu.Unlock()

		queue.Push(&Slice[QueueItem]{
			Key:   key,
			Value: cmd.Item,
		})
		result.Value = uint64(len(d))

		log.Printf("%#v, %d, %#v", s.queues, len(queue.Queue), queue.Queue)

		return result, err
	case "dequeue":

		// prefix := []byte(fmt.Sprintf("%s_", cmd.Group))
		s.mu.Lock()

		var queue *PriorityQueue[QueueItem]
		var ok bool
		if queue, ok = s.queues[cmd.Group]; !ok {
			queue = NewPriorityQueue[QueueItem]()
			s.queues[cmd.Group] = queue
		}
		defer s.mu.Unlock()

		item := queue.Pop()
		if item != nil {
			d, err := item.Encode()
			if err != nil {
				return result, err
			}
			e.Result.Data = d
			result.Data = d
		}
		log.Printf("%#v", s.queues)
		return result, err
	default:
		return result, fmt.Errorf("unknonw cmd type: %s", cmd.Name)
	}

}

// SaveSnapshot saves the current IStateMachine state into a snapshot using the
// specified io.Writer object.
func (s *SMQueue) SaveSnapshot(w io.Writer,
	fc sm.ISnapshotFileCollection, done <-chan struct{}) error {
	// as shown above, the only state that can be saved is the Count variable
	// there is no external file in this IStateMachine example, we thus leave
	// the fc untouched

	s.mu.Lock()
	defer s.mu.Unlock()

	return gob.NewEncoder(w).Encode(&s.queues)
}

// RecoverFromSnapshot recovers the state using the provided snapshot.
func (s *SMQueue) RecoverFromSnapshot(r io.Reader,
	files []sm.SnapshotFile,
	done <-chan struct{}) error {
	// restore the Count variable, that is the only state we maintain in this
	// example, the input files is expected to be empty

	err := gob.NewDecoder(r).Decode(&s.queues)
	if err != nil {
		return err
	}

	return nil
}

// Close closes the IStateMachine instance. There is nothing for us to cleanup
// or release as this is a pure in memory data store. Note that the Close
// method is not guaranteed to be called as node can crash at any time.
func (s *SMQueue) Close() error { return nil }