How to Implement a Custom State Machine in Dragonboat with Snapshot Support

To implement a custom state machine in Dragonboat that supports snapshotting, you must implement either the IStateMachine interface for in-memory state or the IOnDiskStateMachine interface for disk-backed state, providing the three required snapshot lifecycle methods: PrepareSnapshot, SaveSnapshot, and RecoverFromSnapshot.

Dragonboat is a high-performance Go implementation of the Raft consensus protocol. When you implement a custom state machine in Dragonboat, you define how your application state changes in response to committed Raft log entries and how that state is persisted via snapshots. This guide covers the complete implementation path, from choosing the correct interface to registering your factory function with the NodeHost.

Choose the Right State Machine Interface

Dragonboat provides two distinct interface families in the statemachine package. Your choice depends on where your application state lives.

In-Memory State Machines (IStateMachine)

Use IStateMachine when your entire state fits in RAM and you want Dragonboat to manage the snapshotting of that volatile state to disk. This interface is defined in statemachine/rsm.go (lines 42-98) and requires you to implement methods that capture your in-memory structures.

On-Disk State Machines (IOnDiskStateMachine)

Use IOnDiskStateMachine when your state is already persisted to disk (e.g., using BadgerDB, RocksDB, or a custom LSM tree). In this model, the snapshot only needs to capture a small metadata identifier (like the last applied index), while Dragonboat handles the hard-linking of your existing DB files. This interface is defined in statemachine/disk.go (lines 32-132).

Implement the Core State Machine Methods

Both interfaces share the same snapshot lifecycle methods. You must implement these three methods deterministically, as they will be called during snapshot creation and recovery on different nodes in the cluster.

• PrepareSnapshot() (interface{}, error): Returns a lightweight identifier representing the state to be snapshotted (e.g., the current counter value or last applied index).
• SaveSnapshot(ctx interface{}, w io.Writer, files ISnapshotFileCollection, stopc <-chan struct{}) error: Writes the snapshot data to the provided writer. For on-disk state machines, use files.AddFile() to include external DB files.
• RecoverFromSnapshot(r io.Reader, files []SnapshotFile, stopc <-chan struct{}) error: Restores state from the snapshot reader and/or external files.

In-Memory Implementation Example

Here is a complete IStateMachine implementation for a simple counter that supports snapshotting. This example references the interface definitions in statemachine/rsm.go.

package main

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

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

type CounterSM struct {
    value uint64
}

// NewCounterSM is the factory function matching CreateStateMachineFunc.
func NewCounterSM(shardID, replicaID uint64) sm.IStateMachine {
    return &CounterSM{}
}

// Update applies committed Raft entries to the state machine.
func (c *CounterSM) Update(entry sm.Entry) (sm.Result, error) {
    var inc uint64
    if _, err := fmt.Sscanf(string(entry.Cmd), "inc %d", &inc); err != nil {
        return sm.Result{}, err
    }
    c.value += inc
    return sm.Result{Value: c.value}, nil
}

// Lookup handles read-only queries.
func (c *CounterSM) Lookup(arg interface{}) (interface{}, error) {
    return c.value, nil
}

// PrepareSnapshot captures the current state identifier.
func (c *CounterSM) PrepareSnapshot() (interface{}, error) {
    return c.value, nil
}

// SaveSnapshot writes the state to the snapshot writer.
func (c *CounterSM) SaveSnapshot(w io.Writer, collection sm.ISnapshotFileCollection,
    stopc <-chan struct{}) error {
    return binary.Write(w, binary.LittleEndian, c.value)
}

// RecoverFromSnapshot restores state from the snapshot.
func (c *CounterSM) RecoverFromSnapshot(r io.Reader,
    files []sm.SnapshotFile, stopc <-chan struct{}) error {
    var v uint64
    if err := binary.Read(r, binary.LittleEndian, &v); err != nil {
        return err
    }
    c.value = v
    return nil
}

// Close releases resources.
func (c *CounterSM) Close() error { return nil }

On-Disk Implementation Example

For state machines backed by a persistent store like BadgerDB, implement IOnDiskStateMachine as defined in statemachine/disk.go. The key difference is that SaveSnapshot uses ISnapshotFileCollection.AddFile() to include existing DB files rather than streaming data.

type KVStoreSM struct {
    dir  string
    db   *badger.DB
    last uint64
}

// NewKVStoreSM matches CreateOnDiskStateMachineFunc.
func NewKVStoreSM(shardID, replicaID uint64) sm.IOnDiskStateMachine {
    return &KVStoreSM{dir: fmt.Sprintf("kvstore-%d-%d", shardID, replicaID)}
}

// Open initializes the DB and returns the last applied index.
func (k *KVStoreSM) Open(stopc <-chan struct{}) (uint64, error) {
    db, err := badger.Open(badger.DefaultOptions(k.dir))
    if err != nil {
        return 0, err
    }
    k.db = db
    
    // Retrieve last applied index from a metadata key
    var idx uint64
    k.db.View(func(txn *badger.Txn) error {
        item, err := txn.Get([]byte("lastIndex"))
        if err != nil {
            return nil // Treat missing as 0
        }
        return item.Value(func(v []byte) error {
            idx = binary.LittleEndian.Uint64(v)
            return nil
        })
    })
    k.last = idx
    return idx, nil
}

// Update applies a batch of entries atomically.
func (k *KVStoreSM) Update(es []sm.Entry) ([]sm.Entry, error) {
    wb := k.db.NewWriteBatch()
    defer wb.Cancel()
    
    for i, e := range es {
        // Parse command: "set <key> <value>"
        parts := strings.Fields(string(e.Cmd))
        if len(parts) == 3 && parts[0] == "set" {
            wb.Set([]byte(parts[1]), []byte(parts[2]))
        }
        es[i].Result = sm.Result{Value: e.Index}
        k.last = e.Index
    }
    
    // Persist last index
    var buf [8]byte
    binary.LittleEndian.PutUint64(buf[:], k.last)
    wb.Set([]byte("lastIndex"), buf[:])
    
    if err := wb.Flush(); err != nil {
        return nil, err
    }
    return es, nil
}

// Lookup queries a key.
func (k *KVStoreSM) Lookup(arg interface{}) (interface{}, error) {
    key, ok := arg.(string)
    if !ok {
        return nil, fmt.Errorf("invalid key type")
    }
    
    var val []byte
    err := k.db.View(func(txn *badger.Txn) error {
        item, err := txn.Get([]byte(key))
        if err != nil {
            return err
        }
        return item.Value(func(v []byte) error {
            val = append([]byte{}, v...)
            return nil
        })
    })
    return string(val), err
}

// PrepareSnapshot returns the last applied index as the snapshot identifier.
func (k *KVStoreSM) PrepareSnapshot() (interface{}, error) {
    return k.last, nil
}

// SaveSnapshot writes the metadata and registers DB files for hard-linking.
func (k *KVStoreSM) SaveSnapshot(state interface{}, w io.Writer,
    collection sm.ISnapshotFileCollection, stopc <-chan struct{}) error {
    
    // Write the last index (metadata)
    if err := binary.Write(w, binary.LittleEndian, state.(uint64)); err != nil {
        return err
    }
    
    // Register all files in the DB directory for inclusion via hard links
    return filepath.Walk(k.dir, func(path string, info os.FileInfo, err error) error {
        if err != nil || info.IsDir() {
            return err
        }
        // Generate deterministic file ID from relative path
        relPath := strings.TrimPrefix(path, k.dir+"/")
        fileID := xxhash.Sum64String(relPath)
        return collection.AddFile(fileID, path, nil)
    })
}

// RecoverFromSnapshot restores from the metadata and external files.
func (k *KVStoreSM) RecoverFromSnapshot(r io.Reader,
    files []sm.SnapshotFile, stopc <-chan struct{}) error {
    
    // Read the last index
    var last uint64
    if err := binary.Read(r, binary.LittleEndian, &last); err != nil {
        return err
    }
    k.last = last
    
    // Dragonboat handles restoring external files to a temporary location.
    // You should close the current DB, replace files, and reopen via Open().
    if k.db != nil {
        k.db.Close()
        k.db = nil
    }
    // Files are restored by Dragonboat; you typically just need to reopen the DB
    // in the next Open() call or handle the file replacement here.
    return nil
}

// Close shuts down the database.
func (k *KVStoreSM) Close() error {
    if k.db != nil {
        return k.db.Close()
    }
    return nil
}

Register Your Custom State Machine with NodeHost

Once you have implemented the interface and factory function, you must register it with the NodeHost to start a Raft replica. The factory function must match the signature defined in statemachine/rsm.go (lines 25-27) for in-memory state machines or statemachine/disk.go (lines 34-36) for on-disk state machines.

import (
    "github.com/lni/dragonboat/v4"
    "github.com/lni/dragonboat/v4/config"
    sm "github.com/lni/dragonboat/v4/statemachine"
)

func main() {
    // Configure and create the NodeHost
    nhc := config.NodeHostConfig{
        RaftAddress: "localhost:63001",
        // ... other settings
    }
    nh, err := dragonboat.NewNodeHost(nhc)
    if err != nil {
        panic(err)
    }

    // Configuration for the Raft shard
    rc := config.Config{
        ShardID:             1,
        ReplicaID:           1,
        ElectionRTT:         10,
        HeartbeatRTT:        1,
        SnapshotEntries:     10000, // Trigger snapshot every 10k entries
        // Register the factory function
        CreateReplica:       NewCounterSM, // or NewKVStoreSM for on-disk
    }

    // Start the replica
    if err := nh.StartReplica(rc, sm.RegularStateMachine); err != nil {
        panic(err)
    }
}

The StartReplica method in nodehost.go (lines 523-540) consumes your factory and initializes the internal adapter that wraps your state machine implementation.

Configure Snapshot Frequency

Dragonboat triggers snapshots automatically based on the SnapshotEntries or SnapshotThreshold fields in config.Config. When a snapshot is triggered, the NodeHost calls PrepareSnapshot under a mutual-exclusion lock to ensure a consistent state identifier, followed by SaveSnapshot to persist the data.

Tune these configuration values based on your workload characteristics. Frequent snapshots reduce recovery time but increase I/O overhead, while infrequent snapshots minimize I/O but extend the time required to restore a node from log replay.

Summary

To successfully implement a custom state machine in Dragonboat with snapshot support:

• Choose the correct interface: Use IStateMachine from statemachine/rsm.go for in-memory state, or IOnDiskStateMachine from statemachine/disk.go for disk-backed state.
• Implement the snapshot lifecycle: Provide deterministic implementations of PrepareSnapshot, SaveSnapshot, and RecoverFromSnapshot to enable state recovery across replicas.
• Create a factory function: Match the signature CreateStateMachineFunc or CreateOnDiskStateMachineFunc to allow NodeHost to instantiate your state machine.
• Register with NodeHost: Pass your factory to config.Config.CreateReplica and call NodeHost.StartReplica to begin participating in the Raft cluster.
• Configure snapshot triggers: Set SnapshotEntries in config.Config to control how frequently Dragonboat invokes your snapshot methods.

Frequently Asked Questions

What is the difference between IStateMachine and IOnDiskStateMachine?

IStateMachine is designed for state that lives entirely in memory, requiring you to implement SaveSnapshot to serialize your data to the provided io.Writer. IOnDiskStateMachine is designed for state already persisted to disk (like a database), where snapshots only need to capture metadata, and you use ISnapshotFileCollection.AddFile() to include existing DB files via hard links rather than streaming data.

How does Dragonboat handle concurrent snapshots?

Dragonboat guarantees that PrepareSnapshot, SaveSnapshot, and RecoverFromSnapshot are called under mutual exclusion with respect to the Update method. This ensures that the snapshot captures a consistent state without requiring additional locking in your implementation. However, Lookup may still be called concurrently unless you implement your own synchronization.

Can I use external files in my snapshots?

Yes, but only when implementing IOnDiskStateMachine. In the SaveSnapshot method, you receive an ISnapshotFileCollection parameter. You can call collection.AddFile(fileID, path, metadata) to register files from your on-disk database. Dragonboat will efficiently hard-link these files into the snapshot directory rather than copying them. During recovery, RecoverFromSnapshot receives the list of restored files in the files []sm.SnapshotFile parameter.

What happens if RecoverFromSnapshot fails?

If RecoverFromSnapshot returns an error, Dragonboat treats the snapshot as corrupted and will not start the replica. The NodeHost will retry recovery using an earlier snapshot if one exists, or it will attempt to recover by replaying Raft log entries from the beginning. Your implementation must ensure that RecoverFromSnapshot either fully restores the state or returns an error before modifying any internal structures, maintaining atomicity of the recovery process.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}