Auto-commit versus commit-after-process

# /// script
# requires-python = ">=3.10"
# dependencies = ["kafka-python"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Auto-commit against commit-after-process.

10,000 numbered records in a one-partition topic. A consumer "processes"
each by appending its number to a file after a 5 ms sleep. Once with
auto-commit every second, once with a manual commit after every 100 records.
Each run is killed with SIGKILL at a random moment, restarted, and finished;
then the file is checked for missing and duplicated numbers.
Needs the Kafka service from the compose file (port 59092).
"""
import argparse, collections, json, os, random, signal, subprocess, sys, tempfile, time, uuid

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--bootstrap", default="localhost:59092", help="Kafka bootstrap server.")
p.add_argument("--records", type=int, default=10_000, help="Records.")
p.add_argument("--kill-between", type=lambda s: [float(v) for v in s.split(",")], default=[10, 30],
               help="Kill the consumer after a random number of seconds in this range.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--worker", help=argparse.SUPPRESS)
p.add_argument("--topic", help=argparse.SUPPRESS)
p.add_argument("--out", help=argparse.SUPPRESS)
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

if args.worker:                                         # the consumer process
    from kafka import KafkaConsumer
    auto = args.worker == "auto"
    c = KafkaConsumer(args.topic, bootstrap_servers=args.bootstrap, group_id="lab-" + args.topic, enable_auto_commit=auto,
                      auto_commit_interval_ms=1000, auto_offset_reset="earliest", max_poll_records=100,
                      session_timeout_ms=6000, heartbeat_interval_ms=1000,
                      consumer_timeout_ms=30000)          # waits out the dead member's session before giving up
    out = open(args.out, "a")
    for msg in c:
        time.sleep(0.005)
        out.write(msg.value.decode() + "\n"); out.flush()
        if not auto and msg.offset % 100 == 99: c.commit()
    if not auto: c.commit()
    sys.exit(0)

from kafka import KafkaProducer
from kafka.admin import KafkaAdminClient, NewTopic
random.seed(args.seed)
admin = KafkaAdminClient(bootstrap_servers=args.bootstrap)
rows = []
for mode, label in (("auto", "auto-commit every 1 s"), ("manual", "commit after every 100 processed")):
    topic = f"numbers-{uuid.uuid4().hex[:8]}"
    admin.create_topics([NewTopic(topic, 1, 1)]); time.sleep(1)
    pr = KafkaProducer(bootstrap_servers=args.bootstrap)
    for i in range(args.records): pr.send(topic, str(i).encode())
    pr.flush()
    out = os.path.join(tempfile.mkdtemp(), "processed.txt")
    cmd = [sys.executable, __file__, "--worker", mode, "--topic", topic, "--out", out, "--bootstrap", args.bootstrap]
    proc = subprocess.Popen(cmd); time.sleep(random.uniform(*args.kill_between))
    os.kill(proc.pid, signal.SIGKILL); proc.wait()
    subprocess.run(cmd, check=True)                     # restart and finish
    nums = [int(x) for x in open(out).read().split()]
    cnt = collections.Counter(nums)
    missing = args.records - len(cnt); dups = sum(v - 1 for v in cnt.values() if v > 1)
    rows.append({"commit": label, "missing": missing, "duplicated": dups})
    print(f"{label:<34} killed once: {missing:>4} numbers missing, {dups:>4} processed twice")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"runs": rows}}, f)

# A single Kafka broker (KRaft, no ZooKeeper). Clients connect on localhost:59092;
# the PROXY listener (advertised as localhost:59093) is for labs that put a proxy in front of the broker.
services:
  kafka:
    image: apache/kafka:3.9.1
    ports: ["59092:9092", "59094:9094"]
    environment:
      KAFKA_NODE_ID: 1
      KAFKA_PROCESS_ROLES: broker,controller
      KAFKA_LISTENERS: PLAINTEXT://:9092,PROXY://:9094,CONTROLLER://:9093
      KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://localhost:59092,PROXY://localhost:59093
      KAFKA_CONTROLLER_LISTENER_NAMES: CONTROLLER
      KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: CONTROLLER:PLAINTEXT,PLAINTEXT:PLAINTEXT,PROXY:PLAINTEXT
      KAFKA_INTER_BROKER_LISTENER_NAME: PLAINTEXT
      KAFKA_CONTROLLER_QUORUM_VOTERS: 1@localhost:9093
      KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 1
      KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR: 1
      KAFKA_TRANSACTION_STATE_LOG_MIN_ISR: 1
      KAFKA_GROUP_INITIAL_REBALANCE_DELAY_MS: 0
    healthcheck:
      test: ["CMD-SHELL", "/opt/kafka/bin/kafka-broker-api-versions.sh --bootstrap-server localhost:9092 > /dev/null 2>&1 || exit 1"]
      interval: 3s
      retries: 40