Idempotency Layers in Asynchronous Event-Driven Microservices

{
  "headers": {
    "idempotency-key": "a1b2c3d4-e5f6-4a7b-8c9d-0e1f2a3b4c5d",
    "trace-id": "..."
  },
  "payload": {
    "userId": "usr_123",
    "amount": 10000,
    "currency": "USD"
  }
}

// DO NOT USE THIS - FLAWED LOGIC
func handleMessage(key string, payload []byte) error {
    // 1. CHECK
    processed, err := redisClient.Get(ctx, key).Result()
    if err == nil && processed == "COMPLETED" {
        // Already done, acknowledge and return
        return nil 
    }

    // 2. SET (Implicitly, at the end)
    result, err := processBusinessLogic(payload)
    if err != nil {
        return err // NACK and retry
    }

    // 3. SET
    err = redisClient.Set(ctx, key, "COMPLETED", 24 * time.Hour).Err()
    if err != nil {
        // What do we do now? The logic ran but we failed to record it.
        return err
    }

    return nil // ACK message
}

package idempotency

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

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

// StoredState represents the state of an idempotent operation in Redis.
type StoredState struct {
	Status   string `json:"status"`
	Response []byte `json:"response,omitempty"`
}

const (
	StatusStarted   = "STARTED"
	StatusCompleted = "COMPLETED"
)

var ErrRequestInProgress = errors.New("request in progress")

// IdempotencyStore manages the state in Redis.
type IdempotencyStore struct {
	client         *redis.Client
	lockTTL        time.Duration // Time to hold the 'STARTED' lock
	completionTTL  time.Duration // Time to store the final result
}

func NewIdempotencyStore(client *redis.Client, lockTTL, completionTTL time.Duration) *IdempotencyStore {
	return &IdempotencyStore{
		client:        client,
		lockTTL:       lockTTL,
		completionTTL: completionTTL,
	}
}

// MessageHandler defines the function signature for the actual business logic.
type MessageHandler func(ctx context.Context, payload []byte) (response []byte, err error)

// Middleware creates a wrapper around the core business logic handler.
func (s *IdempotencyStore) Middleware(next MessageHandler) MessageHandler {
	return func(ctx context.Context, payload []byte) ([]byte, error) {
		// In a real app, you'd extract this from message headers.
		idempotencyKey, err := extractIdempotencyKeyFromCtx(ctx)
		if err != nil {
			return nil, err // Or handle as a non-idempotent request
		}

		// 1. Attempt to acquire the lock atomically.
		initialState := StoredState{Status: StatusStarted}
		initialStateBytes, _ := json.Marshal(initialState)

		wasSet, err := s.client.SetNX(ctx, idempotencyKey, initialStateBytes, s.lockTTL).Result()
		if err != nil {
			// Redis is down or slow. This is a critical failure path.
			// FAIL CLOSED: Reject the message to prevent duplicate processing.
			return nil, fmt.Errorf("failed to check idempotency key: %w", err)
		}

		if !wasSet {
			// Key already existed, we lost the race. Check its state.
			return s.handleExistingKey(ctx, idempotencyKey)
		}

		// We acquired the lock. Execute the business logic.
		response, businessErr := next(ctx, payload)

		if businessErr != nil {
			// Business logic failed. Release the lock to allow retries.
			s.client.Del(ctx, idempotencyKey)
			return nil, businessErr
		}

		// Business logic succeeded. Store the result permanently.
		finalState := StoredState{Status: StatusCompleted, Response: response}
		finalStateBytes, _ := json.Marshal(finalState)

		err = s.client.Set(ctx, idempotencyKey, finalStateBytes, s.completionTTL).Err()
		if err != nil {
			// CRITICAL EDGE CASE: Logic succeeded but state update failed.
			// The lock will expire and another worker will re-process.
			// This requires careful consideration. Log this as a critical error.
			// Depending on the business logic, you may need a reconciliation process.
			return nil, fmt.Errorf("CRITICAL: failed to save final state for key %s: %w", idempotencyKey, err)
		}

		return response, nil
	}
}

func (s *IdempotencyStore) handleExistingKey(ctx context.Context, key string) ([]byte, error) {
	for i := 0; i < 3; i++ { // Retry loop for in-progress requests
		val, err := s.client.Get(ctx, key).Bytes()
		if err != nil {
			if err == redis.Nil {
				// Key expired between SETNX and GET. This is rare but possible.
				// We can treat this as if we should retry the whole process.
				// For simplicity, we'll return an error to trigger a redelivery.
				return nil, fmt.Errorf("key %s expired during check", key)
			}
			return nil, fmt.Errorf("failed to get existing key: %w", err)
		}

		var state StoredState
		if err := json.Unmarshal(val, &state); err != nil {
			return nil, fmt.Errorf("failed to unmarshal state: %w", err)
		}

		if state.Status == StatusCompleted {
			// Already processed successfully. Return cached response.
			return state.Response, nil
		}

		if state.Status == StatusStarted {
			// Another worker is processing. Wait and retry the check.
			time.Sleep(100 * time.Millisecond)
			continue
		}
	}
	return nil, ErrRequestInProgress // Timed out waiting for other worker
}

// Helper to simulate getting key from context
func extractIdempotencyKeyFromCtx(ctx context.Context) (string, error) {
    key := ctx.Value("idempotency-key")
    if keyStr, ok := key.(string); ok && keyStr != "" {
        return keyStr, nil
    }
    return "", errors.New("idempotency key not found in context")
}

func main() {
    // Setup Redis client
    rdb := redis.NewClient(&redis.Options{
        Addr: "localhost:6379",
    })

    // Setup Idempotency Store
    // 5-minute lock, 24-hour completion storage
    store := NewIdempotencyStore(rdb, 5*time.Minute, 24*time.Hour)

    // This is our actual business logic
    processPaymentHandler := func(ctx context.Context, payload []byte) ([]byte, error) {
        fmt.Println("Executing critical business logic...")
        // ... charge credit card, update database, etc. ...
        time.Sleep(2 * time.Second)
        response := []byte(`{"transactionId": "txn_xyz789", "status": "success"}`)
        fmt.Println("Business logic complete.")
        return response, nil
    }

    // Wrap the business logic with the idempotency middleware
    idempotentHandler := store.Middleware(processPaymentHandler)

    // --- Simulate receiving two identical messages --- 
    ctx1 := context.WithValue(context.Background(), "idempotency-key", "payment-abc-123")
    ctx2 := context.WithValue(context.Background(), "idempotency-key", "payment-abc-123")
    
    var wg sync.WaitGroup
    wg.Add(2)

    // Simulate concurrent consumers
    go func() {
        defer wg.Done()
        fmt.Println("Consumer 1 received message")
        resp, err := idempotentHandler(ctx1, []byte("some payload"))
        if err != nil {
            fmt.Printf("Consumer 1 Error: %v\n", err)
        } else {
            fmt.Printf("Consumer 1 Success: %s\n", string(resp))
        }
    }()

    go func() {
        defer wg.Done()
        time.Sleep(50 * time.Millisecond) // Ensure consumer 1 starts first
        fmt.Println("Consumer 2 received message")
        resp, err := idempotentHandler(ctx2, []byte("some payload"))
        if err != nil {
            fmt.Printf("Consumer 2 Error: %v\n", err)
        } else {
            fmt.Printf("Consumer 2 Success: %s\n", string(resp))
        }
    }()

    wg.Wait()
}

Consumer 1 received message
Executing critical business logic...
Consumer 2 received message
Business logic complete.
Consumer 1 Success: {"transactionId": "txn_xyz789", "status": "success"}
Consumer 2 Success: {"transactionId": "txn_xyz789", "status": "success"}

-- get_or_wait.lua
local key = KEYS[1]
local current_state_json = redis.call('GET', key)

if not current_state_json then
  -- Key doesn't exist, can attempt to acquire lock
  return { "NOT_FOUND" }
end

-- cjson library is available in Redis
local current_state = cjson.decode(current_state_json)

if current_state.status == 'COMPLETED' then
  return { "COMPLETED", current_state_json } 
else
  -- Could be STARTED or another state
  return { "IN_PROGRESS" }
end

    CREATE TABLE idempotency_keys (
        key VARCHAR(255) PRIMARY KEY,
        status VARCHAR(20) NOT NULL,
        response BYTEA,
        locked_at TIMESTAMPTZ,
        created_at TIMESTAMPTZ DEFAULT NOW(),
        updated_at TIMESTAMPTZ DEFAULT NOW()
    );

    INSERT INTO idempotency_keys (key, status, locked_at)
    VALUES ($1, 'STARTED', NOW())
    ON CONFLICT (key) DO NOTHING;

    UPDATE idempotency_keys
    SET status = 'COMPLETED', response = $1, updated_at = NOW()
    WHERE key = $2;