Babuza State Machine Skill

Name: babuza-statemachine
Rating: 62
Author: fanaujie

Package: github.com/fanaujie/babuza/ibabuza

State machine interface implementation for Babuza distributed consensus.

You are an expert at babuza state machine implementation. Help users by:

•Writing code: Implement the ibabuza.BaseStateMachine interface correctly.
•Answering questions: Explain state machine requirements, snapshot handling, and concurrency.

Documentation

Refer to the local files for detailed API:

•./references/statemachine-api.md - State machine interface and implementation guidelines.

Key Patterns

1. Implementing BaseStateMachine Interface

The state machine must implement ibabuza.BaseStateMachine:

import (
    "sync"
    "encoding/json"
    "github.com/fanaujie/babuza/ibabuza"
)

type KVStore struct {
    mu   sync.RWMutex
    data map[string]string
}

// Apply - called serially by Raft, must be deterministic
func (s *KVStore) Apply(e ibabuza.Entry) ibabuza.ApplyResult {
    var cmd Command
    if err := json.Unmarshal(e.Command, &cmd); err != nil {
        return ibabuza.ApplyResult{LogIndex: e.Index, Error: err}
    }
    s.mu.Lock()
    defer s.mu.Unlock()
    s.data[cmd.Key] = cmd.Value
    return ibabuza.ApplyResult{LogIndex: e.Index, Response: "ok"}
}

// Query - may be called concurrently, use proper locking
func (s *KVStore) Query(key any) (any, error) {
    s.mu.RLock()
    defer s.mu.RUnlock()
    return s.data[key.(string)], nil
}

// SaveSnapshot - serialize state to snapshot writer
func (s *KVStore) SaveSnapshot(_ ibabuza.StateMachineSnapshotContext, w ibabuza.StateMachineSnapshotWriter) error {
    s.mu.RLock()
    defer s.mu.RUnlock()

    // Create a file in the snapshot (fileTag identifies the file)
    file, err := w.CreateStateMachineFile("kvdata", 0) // 0 = no compression
    if err != nil {
        return err
    }
    defer file.Close()

    return json.NewEncoder(file).Encode(s.data)
}

// RestoreFromSnapshot - restore state from snapshot reader
func (s *KVStore) RestoreFromSnapshot(r ibabuza.StateMachineSnapshotReader) error {
    s.mu.Lock()
    defer s.mu.Unlock()

    // Open the file from snapshot
    file, _, err := r.Open("kvdata")
    if err != nil {
        return err
    }

    s.data = make(map[string]string)
    return json.NewDecoder(file).Decode(&s.data)
}

// Close - cleanup resources
func (s *KVStore) Close() error { return nil }

2. ApplyResult Structure

type ApplyResult struct {
    LogIndex uint64  // The log index that was applied
    Response any     // Result to return to the proposer (NOT "Result")
    Error    error   // Error if apply failed (field, not method)
}

3. Snapshot Writer/Reader Interfaces

IMPORTANT: These are NOT io.Writer/io.Reader directly!

// StateMachineSnapshotWriter - use CreateStateMachineFile to get io.WriteCloser
type StateMachineSnapshotWriter interface {
    CreateStateMachineFile(fileTag string, compression babuzapb.SnapshotFileCompressionType) (io.WriteCloser, error)
    AddStateMachineFileMetadata(fileTag string, metadata []byte) error
}

// StateMachineSnapshotReader - use Open to get io.Reader
type StateMachineSnapshotReader interface {
    Open(fileTag string) (io.Reader, StateMachineFileDesc, error)
    Metadata() babuzapb.SnapshotMetadata
}

State Machine Requirements

Requirement	Description
Deterministic	`Apply()` must produce same result for same input on all nodes
Thread-safe	`Query()` may be called concurrently with `Apply()`
Serializable	State must be serializable for snapshots
Idempotent	Consider using client sessions for exactly-once semantics

Common Mistakes

•Using Result instead of Response in ApplyResult
•Using snapshot writer/reader as io.Writer/io.Reader - must call CreateStateMachineFile() and Open()
•Non-deterministic Apply() - using time.Now(), random, or external state
•Missing locking in Query() when state is shared with Apply()
•Calling result.Error() as method - it's a field: result.Error

When Writing Code

•Interface: Implement all 5 methods of ibabuza.BaseStateMachine.
•Determinism: Never use time, random, or external state in Apply().
•Locking: Use sync.RWMutex - write lock in Apply(), read lock in Query().
•Snapshot FileTag: Use consistent file tags between SaveSnapshot and RestoreFromSnapshot.
•Error Handling: Return errors in ApplyResult.Error field, check with result.Error != nil.

When Answering Questions

•Determinism: Explain why all nodes must produce identical state from identical log.
•Concurrency: Apply() is serial, Query() may be concurrent.
•Recovery: State is rebuilt from latest snapshot + replayed log entries after restart.