Idempotency Key Management for Asynchronous Distributed Systems

// Example using ioredis in a TypeScript environment
import Redis from 'ioredis';

const redis = new Redis(process.env.REDIS_URL);

// Represents the stored state for an idempotency key
interface IdempotencyRecord {
  status: 'PROCESSING' | 'COMPLETED' | 'FAILED';
  response?: { 
    statusCode: number;
    body: string;
    headers: Record<string, string>;
  };
}

// Lock timeout in seconds. Must be longer than the p99.9 of the operation.
const LOCK_TTL_SECONDS = 3600; 

async function acquireIdempotencyLock(key: string): Promise<'ACQUIRED' | 'DUPLICATE' | 'IN_PROGRESS'> {
  const record: IdempotencyRecord = { status: 'PROCESSING' };
  const result = await redis.set(
    `idempotency:${key}`,
    JSON.stringify(record),
    'EX',
    LOCK_TTL_SECONDS,
    'NX' // SET only if the key does not exist
  );

  if (result === 'OK') {
    return 'ACQUIRED';
  }

  // If we're here, the key already exists. We need to check its state.
  const existingRecordStr = await redis.get(`idempotency:${key}`);
  if (!existingRecordStr) {
      // The key expired between our SET NX and GET. This is a race condition.
      // The client should retry the operation with the same key.
      // We can try to acquire the lock again in a loop, but a simple retry is often sufficient.
      return 'ACQUIRED'; // Or retry logic
  }

  const existingRecord: IdempotencyRecord = JSON.parse(existingRecordStr);
  
  if (existingRecord.status === 'PROCESSING') {
      return 'IN_PROGRESS';
  }
  
  return 'DUPLICATE';
}

async function releaseIdempotencyLock(key: string, successful: boolean, response: any) {
    const record: IdempotencyRecord = {
        status: successful ? 'COMPLETED' : 'FAILED',
        response: response
    };
    // Overwrite the key with the final result. 
    // Use a pipeline to ensure atomicity of setting the value and TTL.
    const pipeline = redis.pipeline();
    pipeline.set(`idempotency:${key}`, JSON.stringify(record), 'EX', LOCK_TTL_SECONDS);
    await pipeline.exec();
}

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

CREATE TABLE idempotency_keys (
    id BIGSERIAL PRIMARY KEY,
    -- The idempotency key provided by the client.
    key VARCHAR(255) NOT NULL,
    -- Scope the key to a specific user or tenant to prevent collisions.
    tenant_id BIGINT NOT NULL REFERENCES tenants(id),
    -- The current state of the operation.
    status idempotency_status NOT NULL DEFAULT 'processing',
    -- The HTTP status code of the final response.
    response_code INT,
    -- The HTTP headers of the final response.
    response_headers JSONB,
    -- The body of the final response.
    response_body BYTEA,
    -- Timestamp for creation and updates, useful for garbage collection.
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    -- This is the critical constraint for ensuring uniqueness.
    UNIQUE (tenant_id, key)
);

-- Index for fast lookups.
CREATE INDEX idx_idempotency_keys_created_at ON idempotency_keys (created_at);

import { Pool, PoolClient } from 'pg';

const pool = new Pool({ connectionString: process.env.DATABASE_URL });

// Define types for clarity
type IdempotencyRecord = {
    id: number;
    key: string;
    tenant_id: number;
    status: 'processing' | 'completed' | 'failed';
    response_code: number | null;
    response_headers: Record<string, any> | null;
    response_body: Buffer | null;
}

type LockResult = 
    | { status: 'acquired'; record: IdempotencyRecord }
    | { status: 'duplicate'; record: IdempotencyRecord }
    | { status: 'in_progress' };

/**
 * Phase 1: Atomically acquire a lock for the idempotency key.
 * This function uses a transaction and pessimistic locking to prevent race conditions.
 */
async function acquireLock(client: PoolClient, key: string, tenantId: number): Promise<LockResult> {
    await client.query('BEGIN');
    try {
        // Take a row-level lock on the potential row. 
        // If the row doesn't exist, Postgres locks the *gap* where the row would be,
        // preventing a concurrent transaction from inserting it.
        const { rows } = await client.query<IdempotencyRecord>(
            'SELECT * FROM idempotency_keys WHERE tenant_id = $1 AND key = $2 FOR UPDATE',
            [tenantId, key]
        );

        if (rows.length > 0) {
            const existingRecord = rows[0];
            await client.query('COMMIT'); // Release the lock immediately
            if (existingRecord.status === 'completed' || existingRecord.status === 'failed') {
                return { status: 'duplicate', record: existingRecord };
            }
            // Another process is working on this. The client should wait and retry.
            return { status: 'in_progress' };
        }

        // The key does not exist. Create it in the 'processing' state.
        const insertResult = await client.query<IdempotencyRecord>(
            'INSERT INTO idempotency_keys (tenant_id, key, status) VALUES ($1, $2, $3) RETURNING *',
            [tenantId, key, 'processing']
        );

        await client.query('COMMIT');
        return { status: 'acquired', record: insertResult.rows[0] };

    } catch (error) {
        await client.query('ROLLBACK');
        // This could be a unique constraint violation if two transactions entered a race
        // before the FOR UPDATE lock was acquired. Retrying the lock acquisition is a safe strategy.
        if (error.code === '23505') { // unique_violation
            return acquireLock(client, key, tenantId); // Recursive retry
        }
        throw error;
    }
}

/**
 * Phase 2: Update the record with the final result after business logic completes.
 * This is done in a separate, short-lived transaction.
 */
async function saveResult(key: string, tenantId: number, code: number, headers: object, body: Buffer) {
    await pool.query(
        `UPDATE idempotency_keys 
         SET status = 'completed', response_code = $3, response_headers = $4, response_body = $5, updated_at = NOW() 
         WHERE tenant_id = $1 AND key = $2`,
        [tenantId, key, code, headers, body]
    );
}

// --- Example Usage in a Webhook Handler ---
async function handleWebhook(request: any) {
    const idempotencyKey = request.headers['idempotency-key'];
    const tenantId = request.user.tenantId;
    const client = await pool.connect();

    try {
        const lock = await acquireLock(client, idempotencyKey, tenantId);

        if (lock.status === 'duplicate') {
            // Serve the stored response
            return {
                statusCode: lock.record.response_code,
                headers: lock.record.response_headers,
                body: lock.record.response_body
            };
        }

        if (lock.status === 'in_progress') {
            // Another request is processing. Tell the client to try again later.
            return { statusCode: 429, body: 'Request in progress' };
        }

        // We have the lock. Execute the business logic.
        const result = await performComplexBusinessLogic(request.body);

        // Save the final result. This is a crucial step.
        await saveResult(idempotencyKey, tenantId, result.statusCode, result.headers, result.body);
        
        return result;

    } catch (error) {
        // Handle business logic errors. Potentially update the key to 'failed'.
        // ...
    } finally {
        client.release();
    }
}

// Example using AWS SDK for JavaScript v3
import { DynamoDBClient, PutItemCommand, GetItemCommand, UpdateItemCommand } from '@aws-sdk/client-dynamodb';

const dynamoClient = new DynamoDBClient({});
const TABLE_NAME = 'IdempotencyKeys';

// Simplified example
async function acquireDynamoLock(key: string, tenantId: string) {
    try {
        // Atomically create the item only if it doesn't exist
        await dynamoClient.send(new PutItemCommand({
            TableName: TABLE_NAME,
            Item: {
                'pk': { S: `T#${tenantId}` },
                'sk': { S: `KEY#${key}` },
                'status': { S: 'processing' },
                'ttl': { N: (Math.floor(Date.now() / 1000) + 3600).toString() }
            },
            ConditionExpression: 'attribute_not_exists(pk) AND attribute_not_exists(sk)'
        }));
        return { status: 'acquired' };
    } catch (error) {
        if (error.name === 'ConditionalCheckFailedException') {
            // The item already exists. We need to check its status.
            const { Item } = await dynamoClient.send(new GetItemCommand({
                TableName: TABLE_NAME,
                Key: {
                    'pk': { S: `T#${tenantId}` },
                    'sk': { S: `KEY#${key}` }
                }
            }));

            if (Item.status.S === 'completed') {
                return { status: 'duplicate', response: JSON.parse(Item.response.S) };
            }
            return { status: 'in_progress' };
        } else {
            throw error;
        }
    }
}

async function saveDynamoResult(key: string, tenantId: string, response: any) {
    await dynamoClient.send(new UpdateItemCommand({
        TableName: TABLE_NAME,
        Key: {
            'pk': { S: `T#${tenantId}` },
            'sk': { S: `KEY#${key}` }
        },
        UpdateExpression: 'SET #status = :status, #response = :response',
        ExpressionAttributeNames: {
            '#status': 'status',
            '#response': 'response'
        },
        ExpressionAttributeValues: {
            ':status': { S: 'completed' },
            ':response': { S: JSON.stringify(response) }
        }
    }));
}

-- This query finds records that have been in 'processing' for more than a configured timeout
-- (e.g., 1 hour) and marks them as 'failed'.
UPDATE idempotency_keys
SET status = 'failed', updated_at = NOW()
WHERE status = 'processing'
  AND updated_at < NOW() - INTERVAL '1 hour';