Kafka share groups with Lambda: set up queue mode

docker run -d --name kafka -p 9092:9092 apache/kafka:4.4.0
docker exec kafka /opt/kafka/bin/kafka-topics.sh --bootstrap-server localhost:9092 \
  --create --topic ride-ended --partitions 1

docker exec -it kafka /opt/kafka/bin/kafka-console-share-consumer.sh \
  --bootstrap-server localhost:9092 --topic ride-ended --group invoice-workers

docker exec -it kafka /opt/kafka/bin/kafka-console-producer.sh \
  --bootstrap-server localhost:9092 --topic ride-ended

docker exec kafka /opt/kafka/bin/kafka-share-groups.sh --bootstrap-server localhost:9092 \
  --describe --group invoice-workers --members

mkdir ride-invoices && cd ride-invoices
npm init -y
npm install aws-cdk-lib constructs @aws-sdk/client-dynamodb @aws-sdk/lib-dynamodb
npm install -D aws-cdk typescript tsx esbuild @types/node @types/aws-lambda
mkdir bin lib src test

// tsconfig.json
{
  "compilerOptions": {
    "target": "ES2022",
    "module": "NodeNext",
    "moduleResolution": "NodeNext",
    "strict": true,
    "esModuleInterop": true,
    "skipLibCheck": true,
    "noEmit": true,
    "types": ["node"]
  },
  "include": ["bin", "lib", "src", "test"]
}

// cdk.json
{
  "app": "npx tsx bin/app.ts"
}

// src/kafka-event.ts
// Minimal types for the Kafka event that Lambda delivers (Stream and Queue mode share the format).

export interface KafkaHeader {
  [name: string]: number[];
}

export interface KafkaRecord {
  topic: string;
  partition: number;
  offset: number;
  timestamp: number;
  timestampType: string;
  key?: string; // base64
  value?: string; // base64
  headers: KafkaHeader[];
}

export interface KafkaEvent {
  eventSource: string; // "SelfManagedKafka"
  bootstrapServers?: string;
  records: Record<string, KafkaRecord[]>; // key: "<topic>-<partition>"
}

export interface RideEnded {
  rideId: string;
  bikeId: string;
  customerEmail: string;
  minutes: number;
}

// Flattens the batch. In Queue mode the order across and within partitions is not guaranteed.
export function flattenRecords(event: KafkaEvent): KafkaRecord[] {
  return Object.values(event.records ?? {}).flat();
}

export function decodeValue(record: KafkaRecord): RideEnded {
  if (!record.value) {
    throw new PoisonMessageError(`Record ${record.topic}/${record.partition}/${record.offset} has no value`);
  }
  let parsed: unknown;
  try {
    parsed = JSON.parse(Buffer.from(record.value, 'base64').toString('utf8'));
  } catch {
    throw new PoisonMessageError(`Record ${record.topic}/${record.partition}/${record.offset} is not valid JSON`);
  }
  const ride = parsed as Partial<RideEnded>;
  if (typeof ride.rideId !== 'string' || typeof ride.minutes !== 'number') {
    throw new PoisonMessageError(`Record ${record.topic}/${record.partition}/${record.offset} is missing rideId or minutes`);
  }
  return ride as RideEnded;
}

// Errors that will never succeed on retry (broken JSON, missing fields).
export class PoisonMessageError extends Error {
  override name = 'PoisonMessageError';
}

// src/idempotency.ts
import { ConditionalCheckFailedException, DynamoDBClient } from '@aws-sdk/client-dynamodb';
import { DeleteCommand, DynamoDBDocumentClient, PutCommand, UpdateCommand } from '@aws-sdk/lib-dynamodb';

export type ClaimResult = 'claimed' | 'already-done' | 'in-progress';

// Small interface so the handler can be tested without AWS.
export interface IdempotencyStore {
  claim(key: string, nowMs: number, lockMs: number): Promise<ClaimResult>;
  complete(key: string, nowMs: number): Promise<void>;
  release(key: string): Promise<void>;
}

const DAY_SECONDS = 24 * 60 * 60;

export class DynamoIdempotencyStore implements IdempotencyStore {
  private readonly doc: DynamoDBDocumentClient;

  constructor(private readonly tableName: string, client = new DynamoDBClient({})) {
    this.doc = DynamoDBDocumentClient.from(client);
  }

  async claim(key: string, nowMs: number, lockMs: number): Promise<ClaimResult> {
    try {
      await this.doc.send(
        new PutCommand({
          TableName: this.tableName,
          Item: {
            pk: key,
            status: 'IN_PROGRESS',
            lockedUntil: nowMs + lockMs,
            expiresAt: Math.floor(nowMs / 1000) + 7 * DAY_SECONDS, // TTL attribute
          },
          // New key, or an abandoned claim whose lock has expired.
          ConditionExpression: 'attribute_not_exists(pk) OR (#s = :inProgress AND lockedUntil < :now)',
          ExpressionAttributeNames: { '#s': 'status' },
          ExpressionAttributeValues: { ':inProgress': 'IN_PROGRESS', ':now': nowMs },
          ReturnValuesOnConditionCheckFailure: 'ALL_OLD',
        }),
      );
      return 'claimed';
    } catch (err) {
      if (err instanceof ConditionalCheckFailedException) {
        return err.Item?.status?.S === 'DONE' ? 'already-done' : 'in-progress';
      }
      throw err;
    }
  }

  async complete(key: string, nowMs: number): Promise<void> {
    await this.doc.send(
      new UpdateCommand({
        TableName: this.tableName,
        Key: { pk: key },
        UpdateExpression: 'SET #s = :done, doneAt = :now REMOVE lockedUntil',
        ExpressionAttributeNames: { '#s': 'status' },
        ExpressionAttributeValues: { ':done': 'DONE', ':now': nowMs },
      }),
    );
  }

  async release(key: string): Promise<void> {
    await this.doc.send(new DeleteCommand({ TableName: this.tableName, Key: { pk: key } }));
  }
}

// In-memory variant for tests and local experiments.
export class MemoryIdempotencyStore implements IdempotencyStore {
  readonly items = new Map<string, { status: 'IN_PROGRESS' | 'DONE'; lockedUntil: number }>();

  async claim(key: string, nowMs: number, lockMs: number): Promise<ClaimResult> {
    const existing = this.items.get(key);
    if (existing?.status === 'DONE') return 'already-done';
    if (existing && existing.lockedUntil >= nowMs) return 'in-progress';
    this.items.set(key, { status: 'IN_PROGRESS', lockedUntil: nowMs + lockMs });
    return 'claimed';
  }

  async complete(key: string): Promise<void> {
    this.items.set(key, { status: 'DONE', lockedUntil: 0 });
  }

  async release(key: string): Promise<void> {
    this.items.delete(key);
  }
}

// src/process-batch.ts
import { decodeValue, flattenRecords, KafkaEvent, PoisonMessageError, RideEnded } from './kafka-event';
import { IdempotencyStore } from './idempotency';

export interface Dependencies {
  store: IdempotencyStore;
  sendInvoice: (ride: RideEnded) => Promise<void>;
  now?: () => number;
  log?: (msg: string, data?: Record<string, unknown>) => void;
}

export interface BatchSummary {
  processed: number;
  skippedDuplicates: number;
  poison: number;
}

// Longer than the function timeout: an abandoned claim only expires once the invocation is surely gone.
export const CLAIM_LOCK_MS = 60_000;

export class RetryLaterError extends Error {
  override name = 'RetryLaterError';
}

export async function processBatch(event: KafkaEvent, deps: Dependencies): Promise<BatchSummary> {
  const now = deps.now ?? Date.now;
  const log = deps.log ?? ((msg, data) => console.log(JSON.stringify({ msg, ...data })));
  const summary: BatchSummary = { processed: 0, skippedDuplicates: 0, poison: 0 };
  let retryNeeded = false;

  for (const record of flattenRecords(event)) {
    let ride: RideEnded;
    try {
      ride = decodeValue(record);
    } catch (err) {
      if (err instanceof PoisonMessageError) {
        // Retrying will not help. Log it and move on so it does not block the batch.
        summary.poison++;
        log('poison message skipped', { reason: err.message });
        continue;
      }
      throw err;
    }

    const key = `ride#${ride.rideId}`; // business key, not the offset
    const claim = await deps.store.claim(key, now(), CLAIM_LOCK_MS);

    if (claim === 'already-done') {
      summary.skippedDuplicates++;
      continue;
    }
    if (claim === 'in-progress') {
      // Another invocation is working on the same ride right now.
      retryNeeded = true;
      continue;
    }

    try {
      await deps.sendInvoice(ride);
      await deps.store.complete(key, now());
      summary.processed++;
    } catch (err) {
      await deps.store.release(key);
      throw err; // the whole batch is retried; finished rides are skipped next time
    }
  }

  if (retryNeeded) {
    throw new RetryLaterError('At least one ride is being processed by another invocation');
  }
  log('batch done', { ...summary });
  return summary;
}

// src/handler.ts
import { DynamoIdempotencyStore } from './idempotency';
import { KafkaEvent, RideEnded } from './kafka-event';
import { processBatch } from './process-batch';

const store = new DynamoIdempotencyStore(process.env.IDEMPOTENCY_TABLE ?? 'missing-table');

// Placeholder for the real work (render the PDF invoice, send the email ...).
async function sendInvoice(ride: RideEnded): Promise<void> {
  console.log(JSON.stringify({ msg: 'invoice sent', rideId: ride.rideId, minutes: ride.minutes }));
}

export const handler = async (event: KafkaEvent) => processBatch(event, { store, sendInvoice });

// lib/ride-invoice-stack.ts
import { Duration, RemovalPolicy, Stack, StackProps } from 'aws-cdk-lib';
import { AttributeType, BillingMode, Table } from 'aws-cdk-lib/aws-dynamodb';
import { CfnEventSourceMapping, Runtime } from 'aws-cdk-lib/aws-lambda';
import { NodejsFunction } from 'aws-cdk-lib/aws-lambda-nodejs';
import { Secret } from 'aws-cdk-lib/aws-secretsmanager';
import { Queue } from 'aws-cdk-lib/aws-sqs';
import { Construct } from 'constructs';

export interface RideInvoiceStackProps extends StackProps {
  bootstrapServers: string[]; // e.g. ["broker1.example.com:9096"]
  kafkaSecretArn: string; // Secrets Manager secret with {"username": "...", "password": "..."}
  topic: string;
}

export class RideInvoiceStack extends Stack {
  constructor(scope: Construct, id: string, props: RideInvoiceStackProps) {
    super(scope, id, props);

    const idempotencyTable = new Table(this, 'IdempotencyTable', {
      partitionKey: { name: 'pk', type: AttributeType.STRING },
      billingMode: BillingMode.PAY_PER_REQUEST,
      timeToLiveAttribute: 'expiresAt',
      removalPolicy: RemovalPolicy.DESTROY, // demo only
    });

    const failedRides = new Queue(this, 'FailedRides', {
      retentionPeriod: Duration.days(14),
    });

    const invoiceFn = new NodejsFunction(this, 'RideInvoiceFn', {
      entry: 'src/handler.ts',
      runtime: Runtime.NODEJS_24_X,
      memorySize: 512,
      // Clearly below the share group's acquisition lock (default 30 s).
      timeout: Duration.seconds(20),
      environment: { IDEMPOTENCY_TABLE: idempotencyTable.tableName },
    });
    idempotencyTable.grantReadWriteData(invoiceFn);
    failedRides.grantSendMessages(invoiceFn);

    const kafkaSecret = Secret.fromSecretCompleteArn(this, 'KafkaSecret', props.kafkaSecretArn);
    kafkaSecret.grantRead(invoiceFn);

    // L1 construct: ConsumptionMode is available directly in CloudFormation.
    new CfnEventSourceMapping(this, 'RideEndedQueueMapping', {
      functionName: invoiceFn.functionName,
      topics: [props.topic],
      selfManagedEventSource: {
        endpoints: { kafkaBootstrapServers: props.bootstrapServers },
      },
      sourceAccessConfigurations: [{ type: 'SASL_SCRAM_512_AUTH', uri: kafkaSecret.secretArn }],
      selfManagedKafkaEventSourceConfig: { consumptionMode: 'Queue' },
      batchSize: 50,
      // No StartingPosition and no MaximumBatchingWindowInSeconds: Queue mode rejects both.
      maximumRetryAttempts: 3,
      bisectBatchOnFunctionError: true,
      destinationConfig: { onFailure: { destination: failedRides.queueArn } },
      provisionedPollerConfig: { minimumPollers: 1, maximumPollers: 10 },
    });
  }
}

// bin/app.ts
import { App } from 'aws-cdk-lib';
import { RideInvoiceStack } from '../lib/ride-invoice-stack';

const app = new App();
const servers = (app.node.tryGetContext('bootstrapServers') as string | undefined) ?? 'broker1.example.com:9096';

new RideInvoiceStack(app, 'RideInvoiceStack', {
  bootstrapServers: servers.split(','),
  kafkaSecretArn:
    (app.node.tryGetContext('kafkaSecretArn') as string | undefined) ??
    'arn:aws:secretsmanager:eu-central-1:111111111111:secret:kafka-ride-invoices-AbCdEf',
  topic: (app.node.tryGetContext('topic') as string | undefined) ?? 'ride-ended',
});

// test/process-batch.test.ts
import assert from 'node:assert/strict';
import { test } from 'node:test';
import { MemoryIdempotencyStore } from '../src/idempotency';
import { KafkaEvent, KafkaRecord, RideEnded } from '../src/kafka-event';
import { processBatch, RetryLaterError } from '../src/process-batch';

function record(partition: number, offset: number, body: unknown): KafkaRecord {
  return {
    topic: 'ride-ended',
    partition,
    offset,
    timestamp: 1_767_000_000_000,
    timestampType: 'CREATE_TIME',
    value: Buffer.from(typeof body === 'string' ? body : JSON.stringify(body)).toString('base64'),
    headers: [],
  };
}

function event(...records: KafkaRecord[]): KafkaEvent {
  const grouped: Record<string, KafkaRecord[]> = {};
  for (const r of records) (grouped[`${r.topic}-${r.partition}`] ??= []).push(r);
  return { eventSource: 'SelfManagedKafka', records: grouped };
}

const ride = (id: string): RideEnded => ({ rideId: id, bikeId: 'b-7', customerEmail: '[email protected]', minutes: 42 });
const silent = () => {};

test('processes records from several partitions', async () => {
  const sent: string[] = [];
  const result = await processBatch(event(record(0, 10, ride('r1')), record(3, 7, ride('r2'))), {
    store: new MemoryIdempotencyStore(),
    sendInvoice: async (r) => void sent.push(r.rideId),
    log: silent,
  });
  assert.deepEqual(result, { processed: 2, skippedDuplicates: 0, poison: 0 });
  assert.deepEqual(sent.sort(), ['r1', 'r2']);
});

test('skips a redelivered ride that is already done', async () => {
  const store = new MemoryIdempotencyStore();
  const sent: string[] = [];
  const deps = { store, sendInvoice: async (r: RideEnded) => void sent.push(r.rideId), log: silent };
  await processBatch(event(record(0, 10, ride('r1'))), deps);
  const second = await processBatch(event(record(0, 10, ride('r1'))), deps);
  assert.equal(second.skippedDuplicates, 1);
  assert.deepEqual(sent, ['r1']);
});

test('skips poison messages instead of blocking the batch', async () => {
  const result = await processBatch(event(record(1, 1, 'not json'), record(1, 2, { bikeId: 'x' }), record(1, 3, ride('r3'))), {
    store: new MemoryIdempotencyStore(),
    sendInvoice: async () => {},
    log: silent,
  });
  assert.deepEqual(result, { processed: 1, skippedDuplicates: 0, poison: 2 });
});

test('releases the claim when sending fails so a retry can process it', async () => {
  const store = new MemoryIdempotencyStore();
  let fail = true;
  const deps = {
    store,
    sendInvoice: async () => {
      if (fail) throw new Error('mail server timeout');
    },
    log: silent,
  };
  await assert.rejects(processBatch(event(record(2, 5, ride('r4'))), deps), /mail server timeout/);
  assert.equal(store.items.has('ride#r4'), false);
  fail = false;
  const retry = await processBatch(event(record(2, 5, ride('r4'))), deps);
  assert.equal(retry.processed, 1);
});

test('asks for a retry when another invocation holds the claim', async () => {
  const store = new MemoryIdempotencyStore();
  await store.claim('ride#r5', Date.now(), 60_000);
  await assert.rejects(
    processBatch(event(record(0, 1, ride('r5')), record(0, 2, ride('r6'))), { store, sendInvoice: async () => {}, log: silent }),
    RetryLaterError,
  );
  assert.equal(store.items.get('ride#r6')?.status, 'DONE');
});

test('takes over an abandoned claim after its lock expired', async () => {
  const store = new MemoryIdempotencyStore();
  await store.claim('ride#r7', 1_000, 60_000);
  const result = await processBatch(event(record(0, 9, ride('r7'))), {
    store,
    sendInvoice: async () => {},
    now: () => 1_000 + 61_000,
    log: silent,
  });
  assert.equal(result.processed, 1);
});

npx tsc --noEmit -p .
npx tsx --test test/*.test.ts

npx cdk bootstrap
npx cdk deploy \
  -c bootstrapServers=broker1.example.com:9096,broker2.example.com:9096 \
  -c kafkaSecretArn=arn:aws:secretsmanager:eu-central-1:111111111111:secret:kafka-ride-invoices-AbCdEf \
  -c topic=ride-ended

aws lambda get-event-source-mapping --uuid <uuid> \
  --query "{state: State, mode: SelfManagedKafkaEventSourceConfig.ConsumptionMode}"