Advanced Idempotency Patterns for Kafka-based Microservices

sequenceDiagram
    participant Producer
    participant Kafka
    participant Consumer
    participant StateStore
    participant BusinessLogic

    Producer->>Kafka: Publish Message (Header: Idempotency-Key: xyz)
    Kafka->>Consumer: Deliver Message
    Consumer->>StateStore: Check if key 'xyz' exists
    alt Key is New
        StateStore-->>Consumer: Not Found
        Consumer->>BusinessLogic: Execute Operation
        BusinessLogic-->>Consumer: Return Result
        Consumer->>StateStore: Store key 'xyz' and Result
        StateStore-->>Consumer: OK
    else Key Exists
        StateStore-->>Consumer: Found, with stored Result
        Consumer-->>Consumer: Skip Business Logic
    end
    Consumer->>Kafka: Commit Offset

package idempotency

import (
	"context"
	"time"

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

// RedisStore implements the idempotency state store using Redis.
type RedisStore struct {
	client *redis.Client
	ttl    time.Duration
}

// NewRedisStore creates a new Redis-backed store.
func NewRedisStore(client *redis.Client, ttl time.Duration) *RedisStore {
	return &RedisStore{client: client, ttl: ttl}
}

var checkAndSetScript = redis.NewScript(`
    -- KEYS[1]: the idempotency key
    -- ARGV[1]: the response payload to store
    -- ARGV[2]: the TTL in seconds
    local current = redis.call('GET', KEYS[1])
    if current then
        return {err = 'conflict'}
    end
    redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
    return {ok = 'stored'}
`)

// CheckAndSet attempts to atomically set the idempotency key.
// It returns the stored response if the key already exists.
func (s *RedisStore) CheckAndSet(ctx context.Context, key string, responsePayload []byte) (existingPayload []byte, isNew bool, err error) {
	// A simple GET is faster if we expect many duplicates.
	existing, err := s.client.Get(ctx, key).Bytes()
	if err == nil {
		return existing, false, nil // Key already exists
	} 
	if err != redis.Nil {
		return nil, false, err // A real error occurred
	}

	// Key does not exist, attempt to set it.
	// Using SETNX is a simpler alternative to Lua for this specific case.
	// SET key value EX ttl NX
	wasSet, err := s.client.SetNX(ctx, key, responsePayload, s.ttl).Result()
	if err != nil {
		return nil, false, err
	}

	if wasSet {
		return nil, true, nil // We successfully set the key
	}

	// If we reach here, a race occurred. Another process set the key between our GET and SETNX.
	// We must read the value again.
	existing, err = s.client.Get(ctx, key).Bytes()
	if err != nil {
		return nil, false, err
	}
	return existing, false, nil
}

CREATE TYPE idempotency_status AS ENUM ('processing', 'completed', 'failed');

CREATE TABLE idempotency_keys (
    -- The idempotency key from the client/producer.
    key_value VARCHAR(255) PRIMARY KEY,

    -- The current state of the operation.
    status idempotency_status NOT NULL,

    -- The consumer group that is processing/processed this key.
    -- Useful for debugging and preventing cross-group interference if needed.
    consumer_group VARCHAR(255) NOT NULL,

    -- The response payload to be returned on subsequent requests.
    -- Storing this allows the consumer to return the exact same response.
    response_payload BYTEA,

    -- Timestamps for lifecycle management and debugging.
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    locked_at TIMESTAMPTZ, -- When the 'processing' lock was acquired.
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

-- Index for efficient cleanup jobs.
CREATE INDEX idx_idempotency_keys_created_at ON idempotency_keys (created_at);

package idempotency

import (
	"context"
	"errors"
	"time"
)

var (
	// ErrConflict indicates that the operation is already being processed or is complete.
	ErrConflict = errors.New("idempotency key conflict")
)

type Status string

const (
	StatusProcessing Status = "processing"
	StatusCompleted  Status = "completed"
	StatusFailed     Status = "failed"
)

// Record represents the state of an idempotent operation.
type Record struct {
	Key             string
	Status          Status
	ResponsePayload []byte
	LockedAt        time.Time
}

// Store defines the interface for the persistent state store.
type Store interface {
	// AcquireLock attempts to lock an idempotency key. 
	// If the key already exists, it returns the existing record and ErrConflict.
	AcquireLock(ctx context.Context, key, consumerGroup string) (*Record, error)

	// SetResult updates the state of a locked key to completed or failed.
	SetResult(ctx context.Context, key string, status Status, responsePayload []byte) error
}

package idempotency

import (
	"context"
	"errors"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
)

// PGStore implements the Store interface using PostgreSQL.
type PGStore struct {
	db *pgxpool.Pool
}

func NewPGStore(db *pgxpool.Pool) *PGStore {
	return &PGStore{db: db}
}

func (s *PGStore) AcquireLock(ctx context.Context, key, consumerGroup string) (*Record, error) {
	var existingRecord Record
	var lockedAt time.Time

	// We use a transaction to ensure atomicity.
	tx, err := s.db.Begin(ctx)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback(ctx) // Rollback is a no-op if tx has been committed.

	// Attempt to insert the key in a 'processing' state. This is our lock.
	insertSQL := `
        INSERT INTO idempotency_keys (key_value, status, consumer_group, locked_at)
        VALUES ($1, 'processing', $2, NOW())
        ON CONFLICT (key_value) DO NOTHING
    `
	cmdTag, err := tx.Exec(ctx, insertSQL, key, consumerGroup)
	if err != nil {
		return nil, err
	}

	// If we successfully inserted a row, we have the lock.
	if cmdTag.RowsAffected() > 0 {
		if err := tx.Commit(ctx); err != nil {
			return nil, err
		}
		return &Record{Key: key, Status: StatusProcessing}, nil
	}

	// If we are here, ON CONFLICT was triggered. The key exists. We must read its state.
	// We use SELECT ... FOR UPDATE to prevent another concurrent transaction from modifying this row
	// until our transaction is complete. This handles the race where two consumers see 'processing'
	// and both decide to wait.
	selectSQL := `
        SELECT status, response_payload, locked_at FROM idempotency_keys 
        WHERE key_value = $1 FOR UPDATE
    `
	err = tx.QueryRow(ctx, selectSQL, key).Scan(&existingRecord.Status, &existingRecord.ResponsePayload, &lockedAt)
	if err != nil {
		return nil, err
	}

	// We can now commit the transaction as we are done with our read lock.
	if err := tx.Commit(ctx); err != nil {
		return nil, err
	}

	// Now, analyze the state we read.
	existingRecord.Key = key
	existingRecord.LockedAt = lockedAt

	// If it's still processing but the lock is old, we might consider it stale.
	// This is an advanced concept: stale lock detection.
	// For now, any 'processing' state is a conflict.
	if existingRecord.Status == StatusProcessing {
		// Optional: Check if locked_at is older than a timeout.
		// if time.Since(lockedAt) > 5*time.Minute { ... handle stale lock ... }
		return &existingRecord, ErrConflict
	}

	// If it's completed or failed, it's also a conflict, but we return the record.
	return &existingRecord, ErrConflict
}

func (s *PGStore) SetResult(ctx context.Context, key string, status Status, responsePayload []byte) error {
	updateSQL := `
        UPDATE idempotency_keys
        SET status = $1, response_payload = $2, updated_at = NOW()
        WHERE key_value = $3
    `
	_, err := s.db.Exec(ctx, updateSQL, status, responsePayload, key)
	return err
}

package main

import (
	"context"
	"fmt"
	"time"

	"example.com/idempotency"
	"github.com/confluentinc/confluent-kafka-go/v2/kafka"
)

// MessageHandler defines the signature for our business logic.
// It returns a byte slice representing the response to be stored.
type MessageHandler func(ctx context.Context, msg *kafka.Message) ([]byte, error)

// IdempotencyMiddleware wraps a MessageHandler to provide idempotency.
func IdempotencyMiddleware(store idempotency.Store, consumerGroup string, next MessageHandler) MessageHandler {
	return func(ctx context.Context, msg *kafka.Message) ([]byte, error) {
		idempotencyKey := getHeader(msg, "Idempotency-Key")
		if idempotencyKey == "" {
			fmt.Println("WARN: No Idempotency-Key found. Processing without idempotency.")
			return next(ctx, msg)
		}

		// 1. Attempt to acquire the lock.
		record, err := store.AcquireLock(ctx, idempotencyKey, consumerGroup)
		if err != nil {
			if err == idempotency.ErrConflict {
				fmt.Printf("Conflict for key %s: status is %s\n", idempotencyKey, record.Status)
				// If already completed, this is a successful duplicate suppression.
				if record.Status == idempotency.StatusCompleted {
					return record.ResponsePayload, nil // Success, do not re-process
				}
				// If it's processing, another consumer has it. We should back off.
				// Returning an error will cause the Kafka library to not commit the offset
				// and the message will be redelivered after a delay.
				return nil, fmt.Errorf("key %s is currently being processed", idempotencyKey)
			}
			// A real database error occurred.
			return nil, err
		}

		// 2. We acquired the lock. Execute the business logic.
		response, err := next(ctx, msg)
		if err != nil {
			// Business logic failed. Mark the key as failed and return the error.
			// We store a nil response payload for failures.
			errUpdate := store.SetResult(ctx, idempotencyKey, idempotency.StatusFailed, nil)
			if errUpdate != nil {
				// If we can't even update the DB, we have a major problem.
				// Log this critical failure.
				fmt.Printf("CRITICAL: Failed to set FAILED status for key %s: %v\n", idempotencyKey, errUpdate)
			}
			return nil, err // Propagate original business logic error
		}

		// 3. Business logic succeeded. Mark as completed and store the response.
		if err := store.SetResult(ctx, idempotencyKey, idempotency.StatusCompleted, response); err != nil {
			// This is also a critical failure. The operation succeeded but we can't mark it as such.
			// This will likely lead to a retry and a duplicate processing attempt.
			// This scenario highlights the need for robust DB error handling and monitoring.
			fmt.Printf("CRITICAL: Failed to set COMPLETED status for key %s: %v\n", idempotencyKey, err)
			return nil, err
		}

		return response, nil
	}
}

func getHeader(msg *kafka.Message, key string) string {
	for _, h := range msg.Headers {
		if h.Key == key {
			return string(h.Value)
		}
	}
	return ""
}

-- This can be run by pg_cron or an external scheduler.
DELETE FROM idempotency_keys
WHERE created_at < NOW() - INTERVAL '48 hours';

-- Create a partitioned table
CREATE TABLE idempotency_keys_partitioned (
    key_value VARCHAR(255) NOT NULL,
    -- ... other columns ...
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
) PARTITION BY RANGE (created_at);

-- The primary key must include the partition key
ALTER TABLE idempotency_keys_partitioned ADD PRIMARY KEY (key_value, created_at);

-- Create partitions automatically or via a script
CREATE TABLE idempotency_keys_2023_10_26 PARTITION OF idempotency_keys_partitioned
    FOR VALUES FROM ('2023-10-26 00:00:00+00') TO ('2023-10-27 00:00:00+00');

graph TD
    subgraph Producer Service
        A[API Request] --> B{DB Transaction Start}
        B --> C[Write to 'orders' table]
        C --> D[Write to 'outbox' table]
        D --> E{Commit Transaction}
    end

    subgraph CDC
        F[Database WAL] -- Debezium --> G[Kafka Producer]
    end

    subgraph Consumer Service
        I[Kafka Consumer] -- Idempotency Middleware --> J[Business Logic]
        J --> K[DB Transaction]
    end

    E --> F
    G --> H[Kafka Topic]
    H --> I