Idempotency Patterns for Kafka Consumers in Distributed Systems

-- The status of an idempotent operation
CREATE TYPE idempotency_status AS ENUM ('started', 'completed', 'failed');

-- Table to store idempotency keys and their state
CREATE TABLE idempotency_keys (
    -- The unique key provided by the producer. Could be scoped per tenant.
    -- For multi-tenant systems, a composite primary key is recommended:
    -- PRIMARY KEY (tenant_id, idempotency_key)
    idempotency_key UUID PRIMARY KEY,

    -- The Kafka partition and offset that triggered this operation.
    -- Useful for tracing and debugging.
    kafka_partition INT NOT NULL,
    kafka_offset BIGINT NOT NULL,

    -- Current status of the operation.
    status idempotency_status NOT NULL DEFAULT 'started',

    -- Timestamp when the lock was first acquired.
    -- Critical for detecting and handling stale locks.
    locked_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    -- The response payload to be returned on subsequent duplicate requests.
    -- Storing this allows the API layer to be idempotent as well.
    response_payload JSONB,

    -- Standard auditing columns.
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    -- Ensure a unique operation per Kafka message offset.
    UNIQUE (kafka_partition, kafka_offset)
);

-- Index for efficient lookup of keys.
CREATE INDEX idx_idempotency_keys_key ON idempotency_keys (idempotency_key);

package main

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

	"github.com/google/uuid"
	_ "github.com/lib/pq"
	"github.com/segmentio/kafka-go"
)

const (
	staleLockTimeout = 5 * time.Minute
)

// Represents the idempotency record in the database
type IdempotencyRecord struct {
	Key             uuid.UUID
	Status          string
	LockedAt        time.Time
	ResponsePayload []byte
}

// Represents the service that processes orders
type OrderService struct {
	DB *sql.DB
}

// ProcessOrderMessage is the core idempotent processing function
func (s *OrderService) ProcessOrderMessage(ctx context.Context, msg kafka.Message) error {
	// 1. Extract Idempotency Key from message header
	idempotencyKeyStr := getHeader(msg.Headers, "idempotency-key")
	if idempotencyKeyStr == "" {
		fmt.Println("Missing idempotency-key header, skipping message")
		return nil // Or move to DLQ
	}
	idempotencyKey, err := uuid.Parse(idempotencyKeyStr)
	if err != nil {
		fmt.Printf("Invalid idempotency-key format: %s\n", idempotencyKeyStr)
		return nil // Or move to DLQ
	}

	// 2. Begin a database transaction
	tx, err := s.DB.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelSerializable})
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback() // Rollback on any error

	// 3. Acquire a pessimistic lock on the idempotency key row
	// This is the most critical step for handling concurrency.
	// `FOR UPDATE` locks the row. If another transaction is holding the lock,
	// this call will block until the lock is released.
	row := tx.QueryRowContext(ctx, 
		`SELECT idempotency_key, status, locked_at, response_payload 
		 FROM idempotency_keys WHERE idempotency_key = $1 FOR UPDATE`, idempotencyKey)

	var record IdempotencyRecord
	err = row.Scan(&record.Key, &record.Status, &record.LockedAt, &record.ResponsePayload)

	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return fmt.Errorf("failed to query idempotency key: %w", err)
	}

	// 4. Handle the different states of the key
	if errors.Is(err, sql.ErrNoRows) { // Case A: First time seeing this key
		fmt.Printf("[%s] New key, starting processing...\n", idempotencyKey)
		_, err := tx.ExecContext(ctx,
			`INSERT INTO idempotency_keys (idempotency_key, kafka_partition, kafka_offset, status) 
			 VALUES ($1, $2, $3, 'started')`,
			idempotencyKey, msg.Partition, msg.Offset)
		if err != nil {
			return fmt.Errorf("failed to insert new idempotency key: %w", err)
		}
	} else { // Case B: Key already exists
		if record.Status == "completed" {
			fmt.Printf("[%s] Duplicate message, operation already completed. Skipping.\n", idempotencyKey)
			// Acknowledge the message without processing
			// Optionally, you could use record.ResponsePayload here
			return tx.Commit()
		}

		if record.Status == "started" && time.Since(record.LockedAt) < staleLockTimeout {
			fmt.Printf("[%s] Operation in progress by another processor. Backing off.\n", idempotencyKey)
			// This indicates a concurrent processing attempt. We can safely abort this one.
			// The other processor holds the lock.
			return errors.New("operation in progress") // This will cause a retry
		}

		// If the lock is stale, this processor takes over.
		fmt.Printf("[%s] Found stale lock, taking over processing.\n", idempotencyKey)
		_, err := tx.ExecContext(ctx, `UPDATE idempotency_keys SET locked_at = NOW() WHERE idempotency_key = $1`, idempotencyKey)
		if err != nil {
			return fmt.Errorf("failed to update stale lock timestamp: %w", err)
		}
	}

	// 5. Execute the core business logic
	// This is where you would charge a card, create an order, etc.
	// For this example, we'll simulate it.
	result, err := s.executeBusinessLogic(ctx, tx, msg.Value)
	if err != nil {
		// On business logic failure, you might want to mark the key as 'failed'
		// to prevent infinite retries for a non-transient error.
		_, updateErr := tx.ExecContext(ctx, `UPDATE idempotency_keys SET status = 'failed' WHERE idempotency_key = $1`, idempotencyKey)
		if updateErr != nil {
			return fmt.Errorf("failed to mark key as failed after business logic error: %w", updateErr)
		}
		return fmt.Errorf("business logic failed: %w", err)
	}

	// 6. Mark the operation as completed and store the result
	resultJSON, _ := json.Marshal(result)
	_, err = tx.ExecContext(ctx, 
		`UPDATE idempotency_keys SET status = 'completed', response_payload = $1, updated_at = NOW() 
		 WHERE idempotency_key = $2`,
		resultJSON, idempotencyKey)
	if err != nil {
		return fmt.Errorf("failed to mark idempotency key as completed: %w", err)
	}

	// 7. Commit the transaction
	// This atomically commits the business logic changes AND the idempotency state update.
	return tx.Commit()
}

func (s *OrderService) executeBusinessLogic(ctx context.Context, tx *sql.Tx, orderData []byte) (map[string]interface{}, error) {
	// Simulate creating an order in the database within the same transaction
	var order map[string]interface{}
	json.Unmarshal(orderData, &order)
	orderID := uuid.New()

	_, err := tx.ExecContext(ctx, `INSERT INTO orders (id, customer_id, amount) VALUES ($1, $2, $3)`,
		orderID, order["customer_id"], order["amount"])
	if err != nil {
		return nil, err
	}

	fmt.Printf("Successfully processed order %s\n", orderID)
	return map[string]interface{}{"order_id": orderID, "status": "created"}, nil
}

func getHeader(headers []kafka.Header, key string) string {
	for _, h := range headers {
		if h.Key == key {
			return string(h.Value)
		}
	}
	return ""
}

import json
import time
import uuid
from typing import Dict, Any

import redis
import psycopg2
from kafka import KafkaConsumer, TopicPartition

# Configuration
REDIS_HOST = 'localhost'
REDIS_PORT = 6379
DB_CONN_STRING = "dbname=orders user=user password=pass host=localhost"
STALE_LOCK_TIMEOUT_SECONDS = 300  # 5 minutes
IDEMPOTENCY_KEY_TTL_SECONDS = 86400 # 24 hours

class HighThroughputOrderProcessor:
    def __init__(self):
        self.redis_client = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, decode_responses=True)
        self.db_conn = psycopg2.connect(DB_CONN_STRING)

    def process_message(self, msg: Any):
        idempotency_key = msg.headers.get('idempotency-key')
        if not idempotency_key:
            print("Missing idempotency-key, skipping.")
            return

        redis_key = f"idempotency:{idempotency_key}"

        # 1. Attempt to acquire a lock atomically using SET NX
        # This is the entry point for all processing attempts.
        lock_payload = json.dumps({"status": "started", "locked_at": time.time()})
        lock_acquired = self.redis_client.set(redis_key, lock_payload, nx=True, ex=STALE_LOCK_TIMEOUT_SECONDS)

        if lock_acquired:
            print(f"[{idempotency_key}] Lock acquired. Starting processing.")
            try:
                # 2. Execute business logic (write to primary DB)
                result = self._execute_business_logic(msg.value)

                # 3. On success, update Redis key to 'completed'
                completion_payload = json.dumps({"status": "completed", "response": result})
                self.redis_client.set(redis_key, completion_payload, ex=IDEMPOTENCY_KEY_TTL_SECONDS)
                print(f"[{idempotency_key}] Processing completed successfully.")

            except Exception as e:
                print(f"[{idempotency_key}] Business logic failed: {e}. Releasing lock.")
                # On failure, we release the lock so another attempt can be made.
                # For non-transient errors, a DLQ strategy is better.
                self.redis_client.delete(redis_key)
                raise # Re-raise to trigger Kafka consumer retry mechanism
        else:
            # Lock was not acquired, another process is working on it.
            self._handle_existing_lock(redis_key, idempotency_key)

    def _handle_existing_lock(self, redis_key: str, idempotency_key: str):
        try:
            raw_payload = self.redis_client.get(redis_key)
            if not raw_payload:
                # The key expired between our SET NX and GET. This is rare but possible.
                # We should retry the whole process.
                print(f"[{idempotency_key}] Key expired mid-check. Retrying will occur.")
                raise Exception("Key expired during check")

            payload = json.loads(raw_payload)

            if payload.get("status") == "completed":
                print(f"[{idempotency_key}] Duplicate message. Operation already complete. Skipping.")
                return # Acknowledge message

            if payload.get("status") == "started":
                locked_at = payload.get("locked_at", 0)
                if (time.time() - locked_at) > STALE_LOCK_TIMEOUT_SECONDS:
                    print(f"[{idempotency_key}] Stale lock detected. Attempting to take over.")
                    # This is a complex recovery path. See discussion below.
                    # For simplicity, we'll just raise an error to retry.
                    # A real implementation needs a more robust recovery strategy.
                    raise Exception("Stale lock detected")
                else:
                    print(f"[{idempotency_key}] Operation in progress. Backing off.")
                    raise Exception("Operation in progress")

        except (json.JSONDecodeError, TypeError):
            print(f"[{idempotency_key}] Invalid payload in Redis. Manual intervention needed.")
            raise Exception("Corrupt idempotency data")

    def _execute_business_logic(self, order_data: bytes) -> Dict[str, Any]:
        order = json.loads(order_data)
        order_id = str(uuid.uuid4())

        # This operation is NOT in a transaction with the Redis update.
        with self.db_conn.cursor() as cur:
            cur.execute(
                "INSERT INTO orders (id, customer_id, amount) VALUES (%s, %s, %s)",
                (order_id, order['customer_id'], order['amount'])
            )
        self.db_conn.commit()

        return {"order_id": order_id, "status": "created"}

    # Inside _handle_existing_lock, when a stale lock is found...
    def recover_from_stale_lock(self, msg: Any) -> bool:
        # 1. Query the primary database to see if the work is already done.
        # This requires a unique identifier from the message payload itself.
        transaction_id = json.loads(msg.value).get('transaction_id')
        with self.db_conn.cursor() as cur:
            cur.execute("SELECT id FROM orders WHERE transaction_id = %s", (transaction_id,))
            existing_order = cur.fetchone()

        if existing_order:
            # The work was done, but the Redis update failed. We fix it now.
            print(f"[{transaction_id}] Recovery: Found existing order. Updating Redis.")
            result = {"order_id": existing_order[0], "status": "created"}
            completion_payload = json.dumps({"status": "completed", "response": result})
            self.redis_client.set(f"idempotency:{msg.headers['idempotency-key']}", completion_payload, ex=IDEMPOTENCY_KEY_TTL_SECONDS)
            return True # Recovery successful, processing is complete.
        else:
            # The work was not done. The previous processor truly failed.
            print(f"[{transaction_id}] Recovery: No existing order found. Proceeding with normal execution.")
            return False # Recovery indicates we should re-run the business logic