Production-Grade Idempotency Layers for Kafka-based Microservices

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

CREATE TABLE idempotency_keys (
    idempotency_key VARCHAR(255) PRIMARY KEY,
    -- The service/consumer group that locked this key
    locking_consumer_id VARCHAR(255),
    
    -- State machine fields
    status idempotency_status NOT NULL DEFAULT 'started',
    
    -- Store the response to return it on subsequent requests
    response_payload JSONB,
    
    -- Timestamps for monitoring and cleanup
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

-- An index is crucial for fast lookups by key
CREATE INDEX idx_idempotency_keys_key ON idempotency_keys(idempotency_key);

package idempotency

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

	"github.com/jackc/pgx/v5/pgconn"
	_ "github.com/jackc/pgx/v5/stdlib"
)

// IdempotencyStatus defines the state of an operation
type IdempotencyStatus string

const (
	StatusStarted    IdempotencyStatus = "started"
	StatusProcessing IdempotencyStatus = "processing"
	StatusCompleted  IdempotencyStatus = "completed"
	StatusFailed     IdempotencyStatus = "failed"
)

// IdempotencyRecord represents the state in the database
type IdempotencyRecord struct {
	Key              string
	Status           IdempotencyStatus
	ResponsePayload  []byte
}

// HandlerFunc is the type for the actual business logic we want to execute idempotently
type HandlerFunc func(ctx context.Context) (interface{}, error)

// ErrDuplicateRequest indicates that the request has already been successfully processed
var ErrDuplicateRequest = errors.New("duplicate request")

// ErrRequestInFlight indicates that the request is currently being processed by another worker
var ErrRequestInFlight = errors.New("request in flight")

// IdempotencyService encapsulates the database connection and logic
type IdempotencyService struct {
	db *sql.DB
}

func NewIdempotencyService(db *sql.DB) *IdempotencyService {
	return &IdempotencyService{db: db}
}

// Handle takes an idempotency key and a handler, and executes the handler idempotently
func (s *IdempotencyService) Handle(ctx context.Context, key string, handler HandlerFunc) (interface{}, error) {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to begin transaction: %w", err)
	}
	// Defer a rollback. If we commit, this is a no-op.
	defer tx.Rollback()

	// Attempt to insert the key first. This is an optimization for the common case (new key).
	_, err = tx.ExecContext(ctx, "INSERT INTO idempotency_keys (idempotency_key, status) VALUES ($1, $2)", key, StatusProcessing)
	if err != nil {
		var pgErr *pgconn.PgError
		// "23505" is the PostgreSQL error code for unique_violation
		if errors.As(err, &pgErr) && pgErr.Code == "23505" {
			// The key already exists, so we need to lock and check its status.
			log.Printf("Key %s already exists, acquiring lock...", key)
			return s.handleExistingKey(ctx, tx, key, handler)
		} else {
			return nil, fmt.Errorf("failed to insert idempotency key: %w", err)
		}
	}

	// If the insert succeeded, we have the lock. Proceed with business logic.
	log.Printf("Key %s is new, processing...", key)
	return s.executeLogic(ctx, tx, key, handler)
}

func (s *IdempotencyService) handleExistingKey(ctx context.Context, tx *sql.Tx, key string, handler HandlerFunc) (interface{}, error) {
	var record IdempotencyRecord
	var responsePayload sql.NullString

	// This is the critical step: acquire a row-level lock.
	// The transaction will block here if another transaction holds the lock.
	// We add a timeout to prevent indefinite waits.
	row := tx.QueryRowContext(ctx, "SELECT status, response_payload FROM idempotency_keys WHERE idempotency_key = $1 FOR UPDATE NOWAIT", key)
	err := row.Scan(&record.Status, &responsePayload)
	if err != nil {
        var pgErr *pgconn.PgError
        // "55P03" is lock_not_available
        if errors.As(err, &pgErr) && pgErr.Code == "55P03" {
            log.Printf("Failed to acquire lock for key %s, another process has it.", key)
            return nil, ErrRequestInFlight
        }
		return nil, fmt.Errorf("failed to scan existing idempotency key: %w", err)
	}

	if responsePayload.Valid {
		record.ResponsePayload = []byte(responsePayload.String)
	}

	switch record.Status {
	case StatusCompleted:
		log.Printf("Key %s already completed. Returning stored response.", key)
		var responseData interface{}
		if err := json.Unmarshal(record.ResponsePayload, &responseData); err != nil {
			return nil, fmt.Errorf("failed to unmarshal stored response: %w", err)
		}
		// We don't need to commit, just returning the old data.
		return responseData, ErrDuplicateRequest
	case StatusProcessing, StatusStarted: 
        // This case is unlikely with FOR UPDATE NOWAIT but could happen if a previous worker crashed.
        // We can treat it as a retryable case.
        log.Printf("Key %s was in processing state. Re-acquiring and executing.", key)
        return s.executeLogic(ctx, tx, key, handler)
	case StatusFailed:
		log.Printf("Key %s was in failed state. Retrying...", key)
		return s.executeLogic(ctx, tx, key, handler)
	default:
		return nil, fmt.Errorf("unknown idempotency status: %s", record.Status)
	}
}

func (s *IdempotencyService) executeLogic(ctx context.Context, tx *sql.Tx, key string, handler HandlerFunc) (interface{}, error) {
	// Execute the actual business logic
	response, err := handler(ctx)
	if err != nil {
		// Business logic failed. Update status to 'failed'.
		_, updateErr := tx.ExecContext(ctx, "UPDATE idempotency_keys SET status = $1, updated_at = NOW() WHERE idempotency_key = $2", StatusFailed, key)
		if updateErr != nil {
			log.Printf("CRITICAL: failed to update idempotency key to FAILED after business logic error: %v", updateErr)
		}
		// We still commit the transaction to record the failure, but return the original business logic error.
        tx.Commit()
		return nil, err
	}

	// Business logic succeeded. Marshal the response and update the status to 'completed'.
	responseBytes, err := json.Marshal(response)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal response: %w", err)
	}

	_, err = tx.ExecContext(ctx, "UPDATE idempotency_keys SET status = $1, response_payload = $2, updated_at = NOW() WHERE idempotency_key = $3", StatusCompleted, responseBytes, key)
	if err != nil {
		return nil, fmt.Errorf("failed to update idempotency key to completed: %w", err)
	}

	// Commit the transaction to release the lock and make the changes permanent.
	if err := tx.Commit(); err != nil {
		return nil, fmt.Errorf("failed to commit transaction: %w", err)
	}

	return response, nil
}

func kafkaMessageHandler(msg kafka.Message) {
    idempotencyKey := string(msg.Headers["idempotency-key"])
    if idempotencyKey == "" {
        log.Println("Missing idempotency key, skipping idempotency check")
        // process message without idempotency or reject it
        return
    }

    // Assume idempotencyService is initialized elsewhere
    response, err := idempotencyService.Handle(context.Background(), idempotencyKey, func(ctx context.Context) (interface{}, error) {
        // Your actual business logic goes here
        // e.g., create an order, update inventory, etc.
        log.Printf("Executing business logic for key %s", idempotencyKey)
        time.Sleep(2 * time.Second) // Simulate work
        return map[string]string{"orderId": "12345", "status": "confirmed"}, nil
    })

    if err != nil {
        if errors.Is(err, idempotency.ErrDuplicateRequest) {
            log.Printf("Handled duplicate message for key: %s", idempotencyKey)
            // Acknowledge the message as processed
            commitKafkaOffset(msg)
            return
        } 
        if errors.Is(err, idempotency.ErrRequestInFlight) {
            log.Printf("Message for key %s is in flight. Will not commit offset, will retry.", idempotencyKey)
            // DO NOT commit offset, let Kafka re-deliver after a timeout
            return
        }
        log.Printf("Business logic failed for key %s: %v", idempotencyKey, err)
        // Decide whether to commit offset based on error type (e.g., move to DLQ)
        return
    }
    
    log.Printf("Successfully processed message for key %s. Response: %v", idempotencyKey, response)
    commitKafkaOffset(msg)
}

import { Redis } from 'ioredis';

const LOCK_TTL_MS = 10000; // Lock expires after 10 seconds
const RESULT_TTL_SEC = 3600; // Stored result expires after 1 hour

export class IdempotencyError extends Error {
  constructor(message: string, public code: 'DUPLICATE' | 'IN_FLIGHT' | 'INTERNAL') {
    super(message);
    this.name = 'IdempotencyError';
  }
}

export class RedisIdempotencyService<TResponse> {
  private redis: Redis;

  constructor(redisClient: Redis) {
    this.redis = redisClient;
  }

  private getLockKey(key: string): string {
    return `idempotency:lock:${key}`;
  }

  private getResultKey(key:string): string {
    return `idempotency:result:${key}`;
  }

  async handle(
    key: string,
    handler: () => Promise<TResponse>
  ): Promise<TResponse> {
    const lockKey = this.getLockKey(key);
    const resultKey = this.getResultKey(key);

    // Step 1: Try to acquire the lock
    const lockAcquired = await this.redis.set(lockKey, 'locked', 'PX', LOCK_TTL_MS, 'NX');

    if (lockAcquired) {
      console.log(`Lock acquired for key: ${key}`);
      try {
        // Step 2: Execute business logic
        const result = await handler();
        const serializedResult = JSON.stringify(result);

        // Step 3: Store result and release lock atomically using a pipeline
        const pipeline = this.redis.pipeline();
        pipeline.set(resultKey, serializedResult, 'EX', RESULT_TTL_SEC);
        pipeline.del(lockKey);
        await pipeline.exec();

        console.log(`Processing completed for key: ${key}`);
        return result;
      } catch (error) {
        // On failure, just release the lock
        console.error(`Business logic failed for key: ${key}. Releasing lock.`, error);
        await this.redis.del(lockKey);
        throw error; // Re-throw original business error
      }
    } else {
      // Lock was not acquired, check for a completed result
      console.log(`Lock not acquired for key: ${key}. Checking for result...`);
      return this.waitForResult(key);
    }
  }

  private async waitForResult(key: string, retries = 5, delay = 200): Promise<TResponse> {
    const resultKey = this.getResultKey(key);

    for (let i = 0; i < retries; i++) {
      const result = await this.redis.get(resultKey);
      if (result) {
        console.log(`Found completed result for key: ${key}`);
        try {
          return JSON.parse(result) as TResponse;
        } catch (e) {
          throw new IdempotencyError('Failed to parse stored result', 'INTERNAL');
        }
      }

      // Wait and try again
      await new Promise(resolve => setTimeout(resolve, delay * Math.pow(2, i)));
    }
    
    throw new IdempotencyError('Request is in flight or processor failed', 'IN_FLIGHT');
  }
}

// Example Usage in a KafkaJS consumer
/*
const redis = new Redis();
const idempotencyService = new RedisIdempotencyService(redis);

consumer.run({ 
  eachMessage: async ({ topic, partition, message }) => {
    const idempotencyKey = message.headers?.['idempotency-key']?.toString();
    if (!idempotencyKey) { /* ... */ /* return; }

    try {
      const result = await idempotencyService.handle(idempotencyKey, async () => {
        // Business logic here
        return { status: 'ok' };
      });
      // On success or duplicate, we are done. Offset is committed automatically.
    } catch (error) {
        if (error instanceof IdempotencyError && error.code === 'IN_FLIGHT') {
            console.warn(`Request for key ${idempotencyKey} is in flight. Will not commit offset.`);
            // CRITICAL: We must throw here to prevent KafkaJS from committing the offset.
            throw error;
        }
        // Handle other business logic errors
        console.error('Processing failed', error);
        throw error; // Let consumer group handle retry/DLQ
    }
  }
})
*/