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

> Implement a custom state machine in Dragonboat with snapshot support. Learn how to integrate IStateMachine or IOnDiskStateMachine and its snapshot lifecycle methods for robust data persistence.

- Repository: [lni/dragonboat](https://github.com/lni/dragonboat)
- Tags: how-to-guide
- Published: 2026-03-06

---

**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`](https://github.com/lni/dragonboat/blob/main/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`](https://github.com/lni/dragonboat/blob/main/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`](https://github.com/lni/dragonboat/blob/main/statemachine/rsm.go).

```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`](https://github.com/lni/dragonboat/blob/main/statemachine/disk.go). The key difference is that `SaveSnapshot` uses `ISnapshotFileCollection.AddFile()` to include existing DB files rather than streaming data.

```go
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`](https://github.com/lni/dragonboat/blob/main/statemachine/rsm.go) (lines 25-27) for in-memory state machines or [`statemachine/disk.go`](https://github.com/lni/dragonboat/blob/main/statemachine/disk.go) (lines 34-36) for on-disk state machines.

```go
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`](https://github.com/lni/dragonboat/blob/main/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`](https://github.com/lni/dragonboat/blob/main/statemachine/rsm.go) for in-memory state, or `IOnDiskStateMachine` from [`statemachine/disk.go`](https://github.com/lni/dragonboat/blob/main/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.