Implementing TCC Distributed Transactions for Redis in Go

Problem

A single MSET command may involve keys that are distributed across multiple Redis nodes using consistent hashing. The operation must be atomic: either all nodes apply the changes or none do.

Two‑Phase Commit (2PC)

2PC splits a write into a Prepare phase and a Commit phase. In the Prepare phase the coordinator (the node that receives the client MSET) sends the transaction to every participant that holds a relevant key. Participants lock the keys, record undo logs, and reply YES if they can proceed or NO /timeout to trigger a rollback. In the Commit phase the coordinator sends a commit request; participants apply the original command, release locks, and acknowledge the coordinator. A Rollback phase is invoked when any participant replies NO or times out.

Issues with a naïve 2PC implementation

Single‑point coordinator : If the coordinator crashes, the transaction cannot progress, risking inconsistency.

Consistency risk : If the coordinator fails after sending Commit requests, some participants may commit while others do not.

Blocking : Participants hold locks from Prepare until Commit, reducing throughput.

Implementation Overview

Source code repository: https://github.com/gofish2020/easyredis. The core files are cluster/router.go and cluster/tcc.go.

Command registration

func init() {
    // Register commands on cluster nodes
    registerClusterRouter("Set", defultFunc)
    registerClusterRouter("Get", defultFunc)
    registerClusterRouter("MSet", mset)

    registerClusterRouter("Prepare", prepareFunc)
    registerClusterRouter("Rollback", rollbackFunc)
    registerClusterRouter("Commit", commitFunc)

    // Direct execution without transaction
    registerClusterRouter("Direct", directFunc)
}

MSet handler

The handler validates arguments, extracts key‑value pairs, groups keys by the node returned by cluster.groupByKeys(keys), and decides between direct execution and a distributed transaction.

func mset(cluster *Cluster, c abstract.Connection, redisCommand [][]byte) protocol.Reply {
    // Basic validation
    if len(redisCommand) < 3 {
        return protocol.NewArgNumErrReply("mset")
    }
    argsNum := len(redisCommand) - 1
    if argsNum%2 != 0 {
        return protocol.NewArgNumErrReply("mset")
    }

    // Extract key‑value pairs
    size := argsNum / 2
    keys := make([]string, 0, size)
    values := make(map[string]string)
    for i := 0; i < size; i++ {
        keys = append(keys, string(redisCommand[2*i+1]))
        values[keys[i]] = string(redisCommand[2*i+2])
    }

    // Map keys to nodes
    ipMap := cluster.groupByKeys(keys)

    // Single‑node case: direct execution
    if len(ipMap) == 1 {
        for ip := range ipMap {
            return cluster.Relay(ip, c, pushCmd(redisCommand, "Direct"))
        }
    }

    // Prepare phase
    var respReply protocol.Reply = protocol.NewOkReply()
    txId := cluster.newTxId()
    rollback := false
    for ip, keys := range ipMap {
        argsGroup := [][]byte{[]byte(txId), []byte("mset")}
        for _, key := range keys {
            argsGroup = append(argsGroup, []byte(key), []byte(values[key]))
        }
        reply := cluster.Relay(ip, c, pushCmd(argsGroup, "Prepare"))
        if protocol.IsErrReply(reply) {
            respReply = reply
            rollback = true
            break
        }
    }

    if rollback {
        rollbackTransaction(cluster, c, txId, ipMap)
    } else {
        _, reply := commitTransaction(cluster, c, txId, ipMap)
        if reply != nil {
            respReply = reply
        }
    }
    return respReply
}

Prepare command

func prepareFunc(cluster *Cluster, conn abstract.Connection, redisCommand [][]byte) protocol.Reply {
    if len(redisCommand) < 3 {
        return protocol.NewArgNumErrReply("prepare")
    }
    txId := string(redisCommand[1])
    // Create transaction object
    tx := NewTransaction(txId, redisCommand[2:], cluster, conn)
    // Store object
    cluster.transactionLock.Lock()
    cluster.transactions[txId] = tx
    cluster.transactionLock.Unlock()
    // Execute prepare logic
    if err := tx.prepare(); err != nil {
        return protocol.NewGenericErrReply(err.Error())
    }
    // Auto‑rollback after timeout to avoid long‑lasting locks
    cluster.delay.Add(maxPrepareTime, genTxKey(txId), func() {
        tx.mu.Lock()
        defer tx.mu.Unlock()
        if tx.status == preparedStatus {
            tx.rollback()
            cluster.transactionLock.Lock()
            delete(cluster.transactions, tx.txId)
            cluster.transactionLock.Unlock()
        }
    })
    return protocol.NewOkReply()
}

Transaction structure

type Transaction struct {
    txId         string
    redisCommand [][]byte
    cluster      *Cluster
    conn         abstract.Connection
    dbIndex      int

    writeKeys  []string
    readKeys   []string
    keysLocked bool
    undoLog    []CmdLine

    status transactionStatus
    mu     *sync.Mutex
}

func NewTransaction(txId string, cmdLine [][]byte, cluster *Cluster, c abstract.Connection) *Transaction {
    return &Transaction{txId: txId, redisCommand: cmdLine, cluster: cluster, conn: c, dbIndex: c.GetDBIndex(), status: createdStatus, mu: &sync.Mutex{}}
}

func (tx *Transaction) prepare() error {
    tx.mu.Lock()
    defer tx.mu.Unlock()
    readKeys, writeKeys := engine.GetRelatedKeys(tx.redisCommand)
    tx.readKeys = readKeys
    tx.writeKeys = writeKeys
    tx.locks()
    tx.undoLog = tx.cluster.engine.GetUndoLogs(tx.dbIndex, tx.redisCommand)
    tx.status = preparedStatus
    return nil
}

Rollback command

func rollbackFunc(cluster *Cluster, conn abstract.Connection, redisCommand [][]byte) protocol.Reply {
    if len(redisCommand) != 2 {
        return protocol.NewArgNumErrReply("rollback")
    }
    cluster.transactionLock.RLock()
    tx, ok := cluster.transactions[string(redisCommand[1])]
    cluster.transactionLock.RUnlock()
    if !ok {
        return protocol.NewIntegerReply(0) // transaction not found
    }
    tx.mu.Lock()
    defer tx.mu.Unlock()
    if err := tx.rollback(); err != nil {
        return protocol.NewGenericErrReply(err.Error())
    }
    // Delay removal of transaction object
    cluster.delay.Add(waitBeforeCleanTx, "", func() {
        cluster.transactionLock.Lock()
        delete(cluster.transactions, tx.txId)
        cluster.transactionLock.Unlock()
    })
    return protocol.NewIntegerReply(1)
}

func (tx *Transaction) rollback() error {
    if tx.status == rolledBackStatus {
        return nil
    }
    tx.locks()
    for _, cmdLine := range tx.undoLog {
        tx.cluster.engine.ExecWithLock(tx.dbIndex, cmdLine)
    }
    tx.unlocks()
    tx.status = rolledBackStatus
    return nil
}

Commit command

func commitFunc(cluster *Cluster, conn abstract.Connection, redisCommand [][]byte) protocol.Reply {
    if len(redisCommand) != 2 {
        return protocol.NewArgNumErrReply("commit")
    }
    cluster.transactionLock.RLock()
    tx, ok := cluster.transactions[string(redisCommand[1])]
    cluster.transactionLock.RUnlock()
    if !ok {
        return protocol.NewIntegerReply(0) // transaction not found
    }
    return tx.commit()
}

func (tx *Transaction) commit() protocol.Reply {
    tx.mu.Lock()
    defer tx.mu.Unlock()
    if tx.status == committedStatus {
        return protocol.NewIntegerReply(0)
    }
    tx.locks()
    reply := tx.cluster.engine.ExecWithLock(tx.dbIndex, tx.redisCommand)
    if protocol.IsErrReply(reply) {
        tx.rollback()
        return reply
    }
    tx.status = committedStatus
    tx.unlocks()
    // Keep transaction object for a short period before cleanup
    tx.cluster.delay.Add(waitBeforeCleanTx, "", func() {
        tx.cluster.transactionLock.Lock()
        delete(tx.cluster.transactions, tx.txId)
        tx.cluster.transactionLock.Unlock()
    })
    return reply
}

State diagrams