Advanced Idempotency Patterns for Kafka Consumers using Redis and Lua

# WARNING: THIS CODE IS FLAWED AND CONTAINS A RACE CONDITION
import redis
import json

# Assume r is a configured Redis client
r = redis.Redis(decode_responses=True)

def process_message_naive(idempotency_key: str, payload: dict):
    redis_key = f"idempotency:v1:{idempotency_key}"

    # 1. Check if already processed
    existing_state_raw = r.get(redis_key)
    if existing_state_raw:
        existing_state = json.loads(existing_state_raw)
        if existing_state['status'] == 'COMPLETED':
            print(f"Key {idempotency_key}: Already completed. Returning stored result.")
            return existing_state['result']
        elif existing_state['status'] == 'PROCESSING':
            print(f"Key {idempotency_key}: Already processing. Skipping.")
            # This is a problem: what if the other process died?
            return None

    # 2. Mark as processing (THE RACE CONDITION IS HERE)
    processing_state = {"status": "PROCESSING"}
    # The gap between GET and SET is the danger zone
    r.set(redis_key, json.dumps(processing_state), ex=300) # 5-minute TTL

    # 3. Execute business logic
    try:
        result = _execute_critical_business_logic(payload)
        
        # 4. Mark as completed
        completed_state = {"status": "COMPLETED", "result": result}
        r.set(redis_key, json.dumps(completed_state), ex=86400) # 24-hour TTL
        return result
    except Exception as e:
        # Handle failure state
        r.delete(redis_key) # Or set to a FAILED state
        raise e

def _execute_critical_business_logic(payload: dict) -> dict:
    # Simulate a database write or API call
    print(f"Executing business logic for payload: {payload}")
    return {"transaction_id": "txn_12345"}

-- File: check_and_set.lua
-- KEYS[1]: The idempotency key (e.g., 'idempotency:v1:some-key')
-- ARGV[1]: The value to set if the key is missing (e.g., '{"status": "PROCESSING"}')
-- ARGV[2]: The TTL for the key in seconds

-- redis.call('get', KEYS[1]) retrieves the current value of the key.
local current_val = redis.call('get', KEYS[1])

-- If the key already exists, return its value.
if current_val then
  return current_val
end

-- If the key does not exist, set it with the provided value and TTL.
-- 'NX' ensures this SET only happens if the key does not exist.
-- 'EX' sets the expiration time.
redis.call('set', KEYS[1], ARGV[1], 'EX', ARGV[2], 'NX')

-- Return nil to indicate that the key was newly set.
return nil

import redis
import json
import time

# Assume r is a configured Redis client
r = redis.Redis(decode_responses=True)

# Load the Lua script
with open('check_and_set.lua', 'r') as f:
    lua_script = f.read()

# Register the script with Redis. This returns a script object that can be called like a function.
# redis-py will use EVALSHA internally, which is more efficient as the script is cached by its SHA1 hash.
check_and_set_script = r.register_script(lua_script)

def process_message_atomic(idempotency_key: str, payload: dict):
    redis_key = f"idempotency:v1:{idempotency_key}"
    processing_ttl = 300  # 5 minutes
    completion_ttl = 86400 # 24 hours

    processing_state_val = json.dumps({"status": "PROCESSING", "timestamp": time.time()})

    # Atomically check and set the key
    # This script returns the existing value if the key exists, or None if it was set.
    existing_state_raw = check_and_set_script(keys=[redis_key], args=[processing_state_val, processing_ttl])

    if existing_state_raw:
        existing_state = json.loads(existing_state_raw)
        status = existing_state.get('status')
        if status == 'COMPLETED':
            print(f"Key {idempotency_key}: Already completed. Returning stored result.")
            return existing_state.get('result')
        elif status == 'PROCESSING':
            print(f"Key {idempotency_key}: Currently processing by another instance. Skipping.")
            # We'll refine this logic later to handle stale processing locks
            return None
        else:
            # Handle other states like FAILED if necessary
            print(f"Key {idempotency_key}: Found in unexpected state '{status}'.")
            return None

    # If we are here, existing_state_raw was None, meaning we successfully set the key to PROCESSING.
    print(f"Key {idempotency_key}: Acquired lock. Processing business logic.")
    try:
        result = _execute_critical_business_logic(payload)
        
        completed_state = {"status": "COMPLETED", "result": result}
        r.set(redis_key, json.dumps(completed_state), ex=completion_ttl)
        print(f"Key {idempotency_key}: Processing complete.")
        return result
    except Exception as e:
        print(f"Key {idempotency_key}: Business logic failed. Releasing lock.")
        # Critical: On failure, we must release the lock so the operation can be retried.
        r.delete(redis_key)
        raise e

# Dummy business logic function
def _execute_critical_business_logic(payload: dict) -> dict:
    print(f"Executing business logic for payload: {payload}")
    time.sleep(2) # Simulate work
    return {"transaction_id": f"txn_{int(time.time())}"}

-- File: acquire_lock.lua
-- KEYS[1]: The idempotency key
-- ARGV[1]: The lock owner token (e.g., 'consumer-1:thread-4:1678886400')
-- ARGV[2]: The processing state value (JSON string)
-- ARGV[3]: The TTL

local current_val_raw = redis.call('get', KEYS[1])

if current_val_raw then
    local current_val = cjson.decode(current_val_raw)
    -- If it's completed, we just return the value.
    if current_val.status == 'COMPLETED' then
        return current_val_raw
    end
    -- If it's processing, we let the caller handle it (e.g., by checking TTL)
    if current_val.status == 'PROCESSING' then
        return current_val_raw
    end
end

-- If key does not exist, or is in a retryable state (e.g. FAILED), acquire the lock.
redis.call('set', KEYS[1], ARGV[2], 'EX', ARGV[3])
return 'LOCKED'

-- File: release_lock.lua
-- KEYS[1]: The idempotency key
-- ARGV[1]: The lock owner token we expect to find
-- ARGV[2]: The new value to set (e.g., the COMPLETED state)
-- ARGV[3]: The new TTL

local current_val_raw = redis.call('get', KEYS[1])
if not current_val_raw then
    return 0 -- Lock doesn't exist
end

local current_val = cjson.decode(current_val_raw)
if current_val.owner == ARGV[1] then
    redis.call('set', KEYS[1], ARGV[2], 'EX', ARGV[3])
    return 1 -- Success
else
    return 0 -- Lock was stolen
end

import redis
import json
import time
import uuid
import os

# A unique identifier for this consumer instance
CONSUMER_ID = f"consumer-{os.getpid()}-{uuid.uuid4()}"

r = redis.Redis(decode_responses=True)

# Assume Lua scripts are loaded and registered as before
# acquire_lock_script = r.register_script(...)
# release_lock_script = r.register_script(...)

# For this example, we'll simulate the script logic in Python for clarity,
# but in production, these should be actual Lua scripts.

def acquire_lock_lua_sim(redis_key, owner_token, processing_state_val, ttl):
    # This is a simplified, NON-ATOMIC simulation of the acquire_lock.lua script
    # In production, use r.evalsha(...) with the actual script
    existing_state_raw = r.get(redis_key)
    if existing_state_raw:
        return existing_state_raw
    
    # The SET NX operation is atomic
    if r.set(redis_key, processing_state_val, ex=ttl, nx=True):
        return 'LOCKED'
    else:
        return r.get(redis_key)

def release_lock_lua_sim(redis_key, owner_token, completed_state_val, ttl):
    # This WATCH/MULTI/EXEC transaction simulates the atomic check-and-set of the release script
    with r.pipeline() as pipe:
        try:
            pipe.watch(redis_key)
            current_state_raw = pipe.get(redis_key)
            if not current_state_raw:
                return 0
            current_state = json.loads(current_state_raw)
            if current_state.get('owner') == owner_token:
                pipe.multi()
                pipe.set(redis_key, completed_state_val, ex=ttl)
                pipe.execute()
                return 1
            return 0
        except redis.exceptions.WatchError:
            # The key was modified by another client after we WATCHed it.
            return 0

def process_message_production(idempotency_key: str, payload: dict):
    redis_key = f"idempotency:v1:{idempotency_key}"
    processing_ttl = 300
    completion_ttl = 86400
    
    # Each attempt gets a unique owner token
    owner_token = f"{CONSUMER_ID}:{uuid.uuid4()}"
    processing_state_val = json.dumps({
        "status": "PROCESSING", 
        "owner": owner_token, 
        "timestamp": time.time()
    })

    # 1. Attempt to acquire the lock
    result = acquire_lock_lua_sim(redis_key, owner_token, processing_state_val, processing_ttl)

    if result == 'LOCKED':
        # We got the lock, proceed
        print(f"Key {idempotency_key} [Owner: {owner_token}]: Lock acquired.")
        pass
    else:
        # Lock is held or operation is complete
        state = json.loads(result)
        if state['status'] == 'COMPLETED':
            print(f"Key {idempotency_key}: Already completed.")
            return state['result']
        
        if state['status'] == 'PROCESSING':
            # This is the key edge case: check if the lock is stale
            lock_timestamp = state.get('timestamp', 0)
            # A stale lock is one that has existed longer than its intended TTL, 
            # even if the Redis key TTL hasn't expired yet (e.g., due to Redis persistence games).
            # A more robust check might involve a heartbeat mechanism.
            if (time.time() - lock_timestamp) > processing_ttl:
                print(f"Key {idempotency_key}: Found stale lock. Overwriting.")
                # This is a complex recovery path. For simplicity, we will skip.
                # A real implementation might try to forcefully acquire the lock.
                return None
            else:
                print(f"Key {idempotency_key}: Actively processing by {state.get('owner')}. Skipping.")
                return None

    # 2. Execute business logic
    try:
        business_result = _execute_critical_business_logic(payload)
        completed_state_val = json.dumps({
            "status": "COMPLETED",
            "result": business_result
        })

        # 3. Atomically release the lock and set the final state
        # This ensures we only write the result if we still own the lock.
        release_success = release_lock_lua_sim(redis_key, owner_token, completed_state_val, completion_ttl)

        if release_success:
            print(f"Key {idempotency_key} [Owner: {owner_token}]: Processing complete, lock released.")
            return business_result
        else:
            # Our lock was stolen! This means our process took too long and another consumer took over.
            # We MUST NOT commit any state. The results of our business logic must be discarded.
            print(f"CRITICAL: Key {idempotency_key} [Owner: {owner_token}]: Lock lost during processing. Aborting.")
            # Here you would need to trigger compensation logic for `_execute_critical_business_logic` if it had side effects.
            return None

    except Exception as e:
        # On failure, we should try to atomically delete the key IF we still own the lock.
        # This allows for a clean retry later.
        print(f"Key {idempotency_key} [Owner: {owner_token}]: Business logic failed.")
        # Implementation of a safe 'delete if owner' script is left as an exercise.
        raise e