Atomic Idempotency Patterns in Event-Driven Systems with Redis & Lua

sequenceDiagram
    participant C as Client
    participant W1 as Worker 1
    participant W2 as Worker 2
    participant DB as Database

    C->>W1: POST /charge (key: key-123)
    C->>W2: POST /charge (key: key-123) [Retry]

    W1->>DB: SELECT result WHERE key='key-123'
    DB-->>W1: (Not Found)

    W2->>DB: SELECT result WHERE key='key-123'
    DB-->>W2: (Not Found)

    note over W1, W2: Both workers believe they are the first!

    W1->>DB: Execute Payment Logic...
    W2->>DB: Execute Payment Logic... (!!) 

    W1->>DB: INSERT result for 'key-123'
    DB-->>W1: OK

    W2->>DB: INSERT result for 'key-123'
    DB-->>W2: Fails (Unique Key Violation)

    note over W2: The operation has already run twice.

// main.go
package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/google/uuid"
)

// Represents the stored result of an idempotent operation.
// Can store either a success payload or an error.
type StoredResponse struct {
	StatusCode int    `json:"statusCode"`
	Body       string `json:"body"`
	Error      string `json:"error,omitempty"`
}

// Custom errors for idempotency logic
var (
	ErrRequestInProgress = errors.New("request in progress")
)

// IdempotencyStore manages the lifecycle of idempotent requests using Redis.
type IdempotencyStore struct {
	client   *redis.Client
	lockTTL  time.Duration // How long the processing lock should be held
	resultTTL time.Duration // How long the final result should be stored
}

func NewIdempotencyStore(client *redis.Client) *IdempotencyStore {
	return &IdempotencyStore{
		client:   client,
		lockTTL:  30 * time.Second,   // A sensible default lock time
		resultTTL: 24 * time.Hour, // Store results for a full day
	}
}

func (s *IdempotencyStore) startedKey(key string) string {
	return fmt.Sprintf("idem:started:%s", key)
}

func (s *IdempotencyStore) resultKey(key string) string {
	return fmt.Sprintf("idem:result:%s", key)
}

// BeginRequest attempts to start processing for an idempotency key.
// It returns a previously stored response if one exists.
// It returns ErrRequestInProgress if another worker holds the lock.
// It returns nil if the lock was successfully acquired.
func (s *IdempotencyStore) BeginRequest(ctx context.Context, key string) (*StoredResponse, error) {
	// 1. Check for a final result first.
	resultData, err := s.client.Get(ctx, s.resultKey(key)).Bytes()
	if err == nil {
		var response StoredResponse
		if jsonErr := json.Unmarshal(resultData, &response); jsonErr == nil {
			fmt.Printf("[IDEM %s] Found cached result.\n", key)
			return &response, nil
		}
		// Handle case where data is corrupted, proceed as if not found
	}
	if err != redis.Nil {
		return nil, fmt.Errorf("redis GET failed: %w", err) // Genuine Redis error
	}

	// 2. Attempt to acquire the lock.
	workerID := uuid.NewString()
	wasSet, err := s.client.SetNX(ctx, s.startedKey(key), workerID, s.lockTTL).Result()
	if err != nil {
		return nil, fmt.Errorf("redis SETNX failed: %w", err)
	}

	if !wasSet {
		fmt.Printf("[IDEM %s] Lock held by another worker.\n", key)
		return nil, ErrRequestInProgress
	}

	fmt.Printf("[IDEM %s] Acquired lock with worker ID %s.\n", key, workerID)
	return nil, nil // Lock acquired, ready to process
}

// CompleteRequest stores the final result of an operation and releases the lock.
func (s *IdempotencyStore) CompleteRequest(ctx context.Context, key string, response StoredResponse) error {
	resultData, err := json.Marshal(response)
	if err != nil {
		return fmt.Errorf("failed to marshal response: %w", err)
	}

	// Use a pipeline to ensure atomicity of setting result and deleting lock
	pipe := s.client.TxPipeline()
	pipe.Set(ctx, s.resultKey(key), resultData, s.resultTTL)
	pipe.Del(ctx, s.startedKey(key))

	_, err = pipe.Exec(ctx)
	if err != nil {
		return fmt.Errorf("redis pipeline failed: %w", err)
	}

	fmt.Printf("[IDEM %s] Stored result and released lock.\n", key)
	return nil
}

// --- Example Usage ---

func main() {
	ctx := context.Background()
	rdb := redis.NewClient(&redis.Options{
		Addr: "localhost:6379",
	})

	store := NewIdempotencyStore(rdb)
	idempotencyKey := "order-45678"

	// Simulate a business operation
	processPayment := func(key string) (*StoredResponse, error) {
		// This is where your core business logic lives
		fmt.Printf("[BIZ %s] Processing payment...\n", key)
		time.Sleep(2 * time.Second) // Simulate work
		fmt.Printf("[BIZ %s] Payment successful.\n", key)
		return &StoredResponse{StatusCode: 200, Body: `{"transactionId": "txn_abc123"}`}, nil
	}

	handleRequest := func(key string) {
		// Begin the idempotent request
		cachedResponse, err := store.BeginRequest(ctx, key)
		if err != nil {
			if err == ErrRequestInProgress {
				fmt.Printf("[HANDLER %s] Request is already in progress. Will not re-process.\n", key)
				return
			}
			fmt.Printf("[HANDLER %s] Error starting request: %v\n", key, err)
			return
		}

		if cachedResponse != nil {
			fmt.Printf("[HANDLER %s] Returning cached response: %+v\n", key, *cachedResponse)
			return
		}

		// Lock acquired, execute the core logic
		response, bizErr := processPayment(key)
		if bizErr != nil {
			// Store the failure result
			failedResponse := StoredResponse{StatusCode: 500, Error: bizErr.Error()}
			store.CompleteRequest(ctx, key, failedResponse)
			return
		}

		// Store the success result
		store.CompleteRequest(ctx, key, *response)
	}

	// Simulate two concurrent requests
	fmt.Println("--- Simulating concurrent requests ---")
	go handleRequest(idempotencyKey)
	go handleRequest(idempotencyKey)

	time.Sleep(5 * time.Second)

	// Simulate a third request after the first has completed
	fmt.Println("\n--- Simulating a later request ---")
	handleRequest(idempotencyKey)

	time.Sleep(2 * time.Second)
}

-- idempotency.lua
-- KEYS[1]: The base idempotency key (e.g., 'order-45678')
-- ARGV[1]: A unique ID for the worker acquiring the lock (e.g., a UUID)
-- ARGV[2]: The TTL for the lock in seconds
-- ARGV[3]: The TTL for the final result in seconds

local result_key = 'idem:result:' .. KEYS[1]
local started_key = 'idem:started:' .. KEYS[1]

-- 1. Check for a finished result first.
local result = redis.call('GET', result_key)
if result then
  -- Operation is already complete, return the cached result.
  return {1, result}
end

-- 2. Check if the operation is already in progress.
local owner = redis.call('GET', started_key)
if owner then
  -- Lock is held. Return status 3 to indicate it's in progress.
  return {3, nil}
end

-- 3. No result and no lock. Attempt to acquire the lock.
redis.call('SET', started_key, ARGV[1], 'EX', ARGV[2])

-- Return status 2 to indicate lock was successfully acquired.
return {2, nil}

// main_lua.go
package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/google/uuid"
)

// StoredResponse and errors are the same as before...

// IdempotencyStore with Lua script support
type IdempotencyStoreLua struct {
	client    *redis.Client
	scriptSHA string
	lockTTL   time.Duration
	resultTTL time.Duration
}

func NewIdempotencyStoreLua(ctx context.Context, client *redis.Client) (*IdempotencyStoreLua, error) {
	luaScript, err := os.ReadFile("idempotency.lua")
	if err != nil {
		return nil, fmt.Errorf("failed to read lua script: %w", err)
	}

	sha, err := client.ScriptLoad(ctx, string(luaScript)).Result()
	if err != nil {
		return nil, fmt.Errorf("failed to load lua script: %w", err)
	}

	fmt.Printf("Loaded idempotency Lua script with SHA: %s\n", sha)

	return &IdempotencyStoreLua{
		client:    client,
		scriptSHA: sha,
		lockTTL:   30 * time.Second,
		resultTTL: 24 * time.Hour,
	}, nil
}

// BeginRequest now uses the atomic Lua script.
func (s *IdempotencyStoreLua) BeginRequest(ctx context.Context, key string) (*StoredResponse, error) {
	workerID := uuid.NewString()

	args := []interface{}{key}
	argv := []interface{}{workerID, s.lockTTL.Seconds(), s.resultTTL.Seconds()}

	res, err := s.client.EvalSha(ctx, s.scriptSHA, args, argv...).Result()
	if err != nil {
		// If the script is not loaded (e.g., Redis restart), load it again.
		if redis.HasErrorPrefix(err, "NOSCRIPT") {
			// This is a simplified fallback. In production, you'd want a more robust
			// script manager to handle this automatically.
			fmt.Println("Script not found, reloading...")
			newStore, loadErr := NewIdempotencyStoreLua(ctx, s.client)
			if loadErr != nil {
				return nil, loadErr
			}
			s.scriptSHA = newStore.scriptSHA
			// Retry the call
			res, err = s.client.EvalSha(ctx, s.scriptSHA, args, argv...).Result()
		}
		if err != nil {
			return nil, fmt.Errorf("failed to execute lua script: %w", err)
		}
	}

	resultSlice, ok := res.([]interface{})
	if !ok || len(resultSlice) != 2 {
		return nil, fmt.Errorf("unexpected lua script result format: %v", res)
	}

	statusCode, _ := resultSlice[0].(int64)

	switch statusCode {
	case 1: // Result was cached
		resultData, _ := resultSlice[1].(string)
		var response StoredResponse
		if jsonErr := json.Unmarshal([]byte(resultData), &response); jsonErr != nil {
			return nil, fmt.Errorf("failed to unmarshal cached response: %w", jsonErr)
		}
		fmt.Printf("[IDEM %s] Found cached result via Lua.\n", key)
		return &response, nil
	case 2: // Lock acquired
		fmt.Printf("[IDEM %s] Acquired lock via Lua with worker ID %s.\n", key, workerID)
		return nil, nil
	case 3: // Lock held by another worker
		fmt.Printf("[IDEM %s] Lock held by another worker (detected by Lua).\n", key)
		return nil, ErrRequestInProgress
	default:
		return nil, fmt.Errorf("unknown status code from lua script: %d", statusCode)
	}
}

// CompleteRequest remains largely the same, but it's still critical.
func (s *IdempotencyStoreLua) CompleteRequest(ctx context.Context, key string, response StoredResponse) error {
	resultData, err := json.Marshal(response)
	if err != nil {
		return fmt.Errorf("failed to marshal response: %w", err)
	}

	// The keys must match the Lua script's logic
	resultKey := fmt.Sprintf("idem:result:%s", key)
	startedKey := fmt.Sprintf("idem:started:%s", key)

	pipe := s.client.TxPipeline()
	pipe.Set(ctx, resultKey, resultData, s.resultTTL)
	pipe.Del(ctx, startedKey)

	_, err = pipe.Exec(ctx)
	if err != nil {
		return fmt.Errorf("redis pipeline failed: %w", err)
	}

	fmt.Printf("[IDEM %s] Stored result and released lock.\n", key)
	return nil
}

// The main function for demonstrating usage would be identical to the previous one,
// just initializing IdempotencyStoreLua instead.

-- complete.lua
-- KEYS[1]: The base idempotency key
-- KEYS[2]: The result key
-- KEYS[3]: The started/lock key
-- ARGV[1]: The worker's unique ID (fencing token)
-- ARGV[2]: The serialized result data
-- ARGV[3]: The result TTL

local current_owner = redis.call('GET', KEYS[3])

-- Only complete the request if we are still the owner of the lock.
if current_owner == ARGV[1] then
  redis.call('SET', KEYS[2], ARGV[2], 'EX', ARGV[3])
  redis.call('DEL', KEYS[3])
  return 1 -- Success
else
  -- We lost the lock. Do not write the result.
  return 0 -- Failure (fenced off)
end