Repository: argoproj-labs/old-argo-dataflow (language: python)

Issue:
Kafka edge-case with pod restart and not offsets committed
I'm writing this up because I think it is an interesting edge-case with Kafka and might be good to document.

* Kafka topic is being pumped with messages using "latest" rather than "earliest".
* Pod starts ups, intending to consume messages.
* Pod is restarted, not having enough time to consume messages from one or more partitions, so does not commit partitions.
* New pod starts up, starting at "latest" because there are not committed messages. It therefore misses messages from one or more partitions.



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 28fa7bf7..8cc7e0f0 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -118,8 +118,12 @@ jobs:
         if: ${{matrix.suite == 'e2e' || matrix.suite == 'examples'}}
       - run: docker pull golang:1.16
       - run: make wait
+      - run: make logs > /tmp/logs &
       - run: make test-${{matrix.suite}}
       - run: git diff
+      - name: cat logs
+        if: ${{ failure() }}
+        run: cat /tmp/logs
       - name: controller logs
         if: ${{ failure() }}
         run: kubectl -n argo-dataflow-system logs deploy/controller-manager -c manager
diff --git a/Makefile b/Makefile
index 3d3afd9b..36b5c05b 100644
--- a/Makefile
+++ b/Makefile
@@ -66,7 +66,7 @@ test-stan-stress:
 test-%:
 	go generate $(shell find ./test/$* -name '*.go')
 	kubectl -n argo-dataflow-system wait pod -l statefulset.kubernetes.io/pod-name --for condition=ready --timeout=2m
-	go test -count 1 -v --tags test ./test/$*
+	go test -failfast -count 1 -v --tags test ./test/$*
 
 pprof:
 	go tool pprof -web http://127.0.0.1:3569/debug/pprof/allocs
@@ -93,6 +93,9 @@ wait:
 	kubectl -n argo-dataflow-system wait deploy --all --for=condition=available --timeout=2m
 	# kubectl wait does not work for statesfulsets, as statefulsets do not have conditions
 	kubectl -n argo-dataflow-system wait pod -l statefulset.kubernetes.io/pod-name --for condition=ready
+$(GOBIN)/stern:
+	curl -Lo $(GOBIN)/stern https://github.com/wercker/stern/releases/download/1.11.0/stern_`uname -s|tr '[:upper:]' '[:lower:]'`_amd64
+	chmod +x $(GOBIN)/stern
 logs: $(GOBIN)/stern
 	stern -n argo-dataflow-system --tail=3 -l dataflow.argoproj.io/step-name .
 
diff --git a/runner/sidecar/sink/kafka/kafka.go b/runner/sidecar/sink/kafka/kafka.go
index 413eb36d..cae778dd 100644
--- a/runner/sidecar/sink/kafka/kafka.go
+++ b/runner/sidecar/sink/kafka/kafka.go
@@ -35,6 +35,15 @@ func New(ctx context.Context, sinkName string, secretInterface corev1.SecretInte
 	if x.MaxMessageBytes > 0 {
 		config["message.max.bytes"] = x.Kafka.MaxMessageBytes
 	}
+	// https://docs.confluent.io/cloud/current/client-apps/optimizing/throughput.html
+	config["batch.size"] = 100000
+	if x.Async {
+		config["linger.ms"] = 50 // Alias for queue.buffering.max.ms
+	} else {
+		config["linger.ms"] = 0
+	}
+	config["compression.type"] = "lz4"
+	config["request.required.acks"] = "all"
 	// https://github.com/confluentinc/confluent-kafka-go/blob/master/examples/producer_example/producer_example.go
 	producer, err := kafka.NewProducer(&config)
 	if err != nil {
diff --git a/runner/sidecar/source/kafka/kafka.go b/runner/sidecar/source/kafka/kafka.go
index c8f7bbc6..7d9fa478 100644
--- a/runner/sidecar/source/kafka/kafka.go
+++ b/runner/sidecar/source/kafka/kafka.go
@@ -3,7 +3,9 @@ package kafka
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"math"
 	"sync"
 	"time"
 
@@ -31,6 +33,11 @@ type kafkaSource struct {
 	totalLag   int64
 }
 
+const (
+	seconds            = 1000
+	pendingUnavailable = math.MinInt32
+)
+
 func New(ctx context.Context, secretInterface corev1.SecretInterface, mntr monitor.Interface, consumerGroupID, sourceName, sourceURN string, replica int, x dfv1.KafkaSource, process source.Process) (source.Interface, error) {
 	logger := sharedutil.NewLogger().WithValues("source", sourceName)
 	config, err := sharedkafka.GetConfig(ctx, secretInterface, x.KafkaConfig)
@@ -39,6 +46,8 @@ func New(ctx context.Context, secretInterface corev1.SecretInterface, mntr monit
 	}
 	config["group.id"] = consumerGroupID
 	config["group.instance.id"] = fmt.Sprintf("%s/%d", consumerGroupID, replica)
+	config["heartbeat.interval.ms"] = 3 * seconds
+	config["socket.keepalive.enable"] = true
 	config["enable.auto.commit"] = false
 	config["enable.auto.offset.store"] = false
 	if x.StartOffset == "First" {
@@ -46,7 +55,10 @@ func New(ctx context.Context, secretInterface corev1.SecretInterface, mntr monit
 	} else {
 		config["auto.offset.reset"] = "latest"
 	}
-	config["statistics.interval.ms"] = 5 * 1000
+	config["statistics.interval.ms"] = 5 * seconds
+	// https://docs.confluent.io/cloud/current/client-apps/optimizing/throughput.html
+	// config["fetch.min.bytes"] = 100000
+	// config["fetch.max.wait.ms"] = seconds / 2
 	logger.Info("Kafka config", "config", sharedutil.MustJSON(sharedkafka.RedactConfigMap(config)))
 	// https://github.com/confluentinc/confluent-kafka-go/blob/master/examples/consumer_example/consumer_example.go
 	consumer, err := kafka.NewConsumer(&config)
@@ -63,6 +75,7 @@ func New(ctx context.Context, secretInterface corev1.SecretInterface, mntr monit
 		channels:   map[int32]chan *kafka.Message{}, // partition -> messages
 		wg:         &sync.WaitGroup{},
 		process:    process,
+		totalLag:   pendingUnavailable,
 	}
 
 	if err = consumer.Subscribe(x.Topic, func(consumer *kafka.Consumer, event kafka.Event) error {
@@ -111,22 +124,6 @@ func (s *kafkaSource) revokedPartition(partition int32) {
 	}
 }
 
-type Stats struct {
-	Topics map[string]struct {
-		Partitions map[string]struct {
-			ConsumerLag int64 `json:"consumer_lag"`
-		} `json:"partitions"`
-	} `json:"topics"`
-}
-
-func (s Stats) totalLag(topic string) int64 {
-	var totalLag int64
-	for _, p := range s.Topics[topic].Partitions {
-		totalLag += p.ConsumerLag
-	}
-	return totalLag
-}
-
 func (s *kafkaSource) startPollLoop(ctx context.Context) {
 	s.logger.Info("starting poll loop")
 	for {
@@ -182,7 +179,9 @@ func (s *kafkaSource) Close() error {
 }
 
 func (s *kafkaSource) GetPending(context.Context) (uint64, error) {
-	if s.totalLag >= 0 {
+	if s.totalLag == pendingUnavailable {
+		return 0, source.ErrPendingUnavailable
+	} else if s.totalLag >= 0 {
 		return uint64(s.totalLag), nil
 	} else {
 		return 0, nil
@@ -208,18 +207,29 @@ func (s *kafkaSource) consumePartition(ctx context.Context, partition int32) {
 	logger := s.logger.WithValues("partition", partition)
 	logger.Info("consuming partition")
 	s.wg.Add(1)
+	var firstCommittedOffset, lastCommittedOffset int64 = -1, -1
 	defer func() {
-		logger.Info("done consuming partition")
+		logger.Info("done consuming partition", "firstCommittedOffset", firstCommittedOffset, "lastCommittedOffset", lastCommittedOffset)
 		s.wg.Done()
 	}()
 	for msg := range s.channels[partition] {
 		offset := int64(msg.TopicPartition.Offset)
 		logger := logger.WithValues("offset", offset)
 		if err := s.processMessage(ctx, msg); err != nil {
-			logger.Error(err, "failed to process message")
+			if errors.Is(err, context.Canceled) {
+				logger.Info("failed to process message", "err", err.Error())
+			} else {
+				logger.Error(err, "failed to process message")
+			}
 		} else {
 			if _, err := s.consumer.CommitMessage(msg); err != nil {
 				logger.Error(err, "failed to commit message")
+			} else {
+				if firstCommittedOffset == -1 {
+					firstCommittedOffset = offset
+					logger.Info("offset", "firstCommittedOffset", firstCommittedOffset)
+				}
+				lastCommittedOffset = offset
 			}
 		}
 	}
diff --git a/runner/sidecar/source/kafka/stats.go b/runner/sidecar/source/kafka/stats.go
new file mode 100644
index 00000000..9ec09b05
--- /dev/null
+++ b/runner/sidecar/source/kafka/stats.go
@@ -0,0 +1,17 @@
+package kafka
+
+type Stats struct {
+	Topics map[string]struct {
+		Partitions map[string]struct {
+			ConsumerLag int64 `json:"consumer_lag"`
+		} `json:"partitions"`
+	} `json:"topics"`
+}
+
+func (s Stats) totalLag(topic string) int64 {
+	var totalLag int64
+	for _, p := range s.Topics[topic].Partitions {
+		totalLag += p.ConsumerLag
+	}
+	return totalLag
+}
diff --git a/runner/sidecar/source/source.go b/runner/sidecar/source/source.go
index ca6af781..20568317 100644
--- a/runner/sidecar/source/source.go
+++ b/runner/sidecar/source/source.go
@@ -2,6 +2,7 @@ package source
 
 import (
 	"context"
+	"errors"
 	"io"
 )
 
@@ -11,7 +12,11 @@ type Interface interface {
 
 type Process func(ctx context.Context, msg []byte) error
 
+var ErrPendingUnavailable = errors.New("pending not available")
+
 type HasPending interface {
 	Interface
+	// GetPending returns the number of pending messages.
+	// It may return ErrPendingUnavailable if this is not available yet.
 	GetPending(ctx context.Context) (uint64, error)
 }
diff --git a/runner/sidecar/sources.go b/runner/sidecar/sources.go
index e7fd2b8f..3a5a8acf 100644
--- a/runner/sidecar/sources.go
+++ b/runner/sidecar/sources.go
@@ -2,6 +2,7 @@ package sidecar
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"
 
@@ -91,7 +92,6 @@ func connectSources(ctx context.Context, process func(context.Context, []byte) e
 			totalBytesCounter.WithLabelValues(sourceName, fmt.Sprint(replica)).Add(float64(len(msg)))
 
 			meta, err := dfv1.MetaFromContext(ctx)
-
 			if err != nil {
 				return fmt.Errorf("could not send message: %w", err)
 			}
@@ -190,7 +190,11 @@ func connectSources(ctx context.Context, process func(context.Context, []byte) e
 			logger.Info("starting pending loop", "source", sourceName, "updateInterval", updateInterval.String())
 			go wait.JitterUntilWithContext(ctx, func(ctx context.Context) {
 				if pending, err := x.GetPending(ctx); err != nil {
-					logger.Error(err, "failed to get pending", "source", sourceName)
+					if errors.Is(err, source.ErrPendingUnavailable) {
+						logger.Info("failed to get pending", "source", sourceName, "err", err.Error())
+					} else {
+						logger.Error(err, "failed to get pending", "source", sourceName)
+					}
 				} else {
 					logger.Info("got pending", "source", sourceName, "pending", pending)
 					pendingGauge.WithLabelValues(sourceName).Set(float64(pending))


Test patch:
diff --git a/test/db-e2e/db_source_test.go b/test/db-e2e/db_source_test.go
index 968c0fb5..34a896b7 100644
--- a/test/db-e2e/db_source_test.go
+++ b/test/db-e2e/db_source_test.go
@@ -108,7 +108,6 @@ func TestDBSource(t *testing.T) {
 
 	WaitForSunkMessages()
 	WaitForTotalSunkMessages(3)
-	WaitForProcessLatencySeconds(0)
 
 	DeletePipelines()
 	WaitForPodsToBeDeleted()
diff --git a/test/fixtures.go b/test/fixtures.go
index c6143c1b..272e140e 100644
--- a/test/fixtures.go
+++ b/test/fixtures.go
@@ -4,7 +4,6 @@ package test
 
 import (
 	"log"
-	"runtime"
 	"testing"
 
 	"k8s.io/client-go/dynamic"
@@ -46,21 +45,14 @@ func Setup(t *testing.T) (teardown func()) {
 		log.Printf("\n")
 		stopTestAPIPortForward()
 		log.Printf("\n")
-		r := recover()                       // tests should panic on error, we recover so we can run other tests
-		tailLogs := runtime.GOOS != "darwin" // we're probably running on CI
+		r := recover() // tests should panic on error, we recover so we can run other tests
 		if r != nil {
 			log.Printf("❌ FAIL: %s %v\n", t.Name(), r)
 			log.Printf("\n")
-			if tailLogs {
-				TailLogs()
-			}
 			t.Fail()
 		} else if t.Failed() {
 			log.Printf("❌ FAIL: %s\n", t.Name())
-			if tailLogs {
-				log.Printf("\n")
-			}
-			TailLogs()
+			log.Printf("\n")
 		} else {
 			log.Printf("✅ PASS: %s\n", t.Name())
 			log.Printf("\n")
diff --git a/test/kafka-e2e/kafka_test.go b/test/kafka-e2e/kafka_test.go
index 8561472f..ac4aaef3 100644
--- a/test/kafka-e2e/kafka_test.go
+++ b/test/kafka-e2e/kafka_test.go
@@ -21,7 +21,7 @@ func TestKafka(t *testing.T) {
 	topic := CreateKafkaTopic()
 	sinkTopic := CreateKafkaTopic()
 
-	name := CreatePipeline(Pipeline{
+	CreatePipeline(Pipeline{
 		ObjectMeta: metav1.ObjectMeta{GenerateName: "kafka-"},
 		Spec: PipelineSpec{
 			Steps: []StepSpec{{
@@ -37,7 +37,7 @@ func TestKafka(t *testing.T) {
 
 	PumpKafkaTopic(topic, 17)
 
-	defer StartPortForward(name + "-main-0")()
+	defer StartPortForward()()
 	WaitForSunkMessages()
 
 	WaitForTotalSourceMessages(17)
@@ -53,7 +53,7 @@ func TestKafkaAsync(t *testing.T) {
 	topic := CreateKafkaTopic()
 	sinkTopic := CreateKafkaTopic()
 
-	name := CreatePipeline(Pipeline{
+	CreatePipeline(Pipeline{
 		ObjectMeta: metav1.ObjectMeta{GenerateName: "kafka-"},
 		Spec: PipelineSpec{
 			Steps: []StepSpec{{
@@ -67,7 +67,7 @@ func TestKafkaAsync(t *testing.T) {
 	WaitForPipeline()
 	WaitForPod()
 
-	defer StartPortForward(name + "-main-0")()
+	defer StartPortForward()()
 
 	PumpKafkaTopic(topic, 17)
 
diff --git a/test/kafka-fmea/kafka_fmea_test.go b/test/kafka-fmea/kafka_fmea_test.go
index a2a7d67e..1c4a65e7 100644
--- a/test/kafka-fmea/kafka_fmea_test.go
+++ b/test/kafka-fmea/kafka_fmea_test.go
@@ -22,10 +22,8 @@ func TestKafkaFMEA_PodDeletedDisruption(t *testing.T) {
 	topic := CreateKafkaTopic()
 	sinkTopic := CreateKafkaTopic()
 
-	time.Sleep(10*time.Second)
-
 	name := CreatePipeline(Pipeline{
-		ObjectMeta: metav1.ObjectMeta{GenerateName: "kafka-"},
+		ObjectMeta: metav1.ObjectMeta{GenerateName: "kafka-pod-deleted-disruption-"},
 		Spec: PipelineSpec{
 			Steps: []StepSpec{{
 				Name:    "main",
@@ -42,11 +40,15 @@ func TestKafkaFMEA_PodDeletedDisruption(t *testing.T) {
 	n := 500 * 15
 	go PumpKafkaTopic(topic, n)
 
+	stopPortForward := StartPortForward()
+	WaitForNSunkMessages(1000)
+	stopPortForward()
+
 	DeletePod(name + "-main-0") // delete the pod to see that we recover and continue to process messages
-	WaitForPod(name + "-main-0")
+	WaitForPod()
 
 	ExpectKafkaTopicCount(sinkTopic, n, 2*time.Minute)
-	defer StartPortForward(name + "-main-0")()
+	defer StartPortForward()()
 	WaitForNoErrors()
 }
 
@@ -58,29 +60,32 @@ func TestKafkaFMEA_KafkaServiceDisruption(t *testing.T) {
 	topic := CreateKafkaTopic()
 	sinkTopic := CreateKafkaTopic()
 	CreatePipeline(Pipeline{
-		ObjectMeta: metav1.ObjectMeta{Name: "kafka-"},
+		ObjectMeta: metav1.ObjectMeta{GenerateName: "kafka-service-disruption-"},
 		Spec: PipelineSpec{
 			Steps: []StepSpec{{
 				Name:    "main",
 				Cat:     &Cat{},
-				Sources: []Source{{Kafka: &KafkaSource{Kafka: Kafka{Topic: topic}}}},
+				Sources: []Source{{Kafka: &KafkaSource{StartOffset: "First", Kafka: Kafka{Topic: topic}}}},
 				Sinks:   []Sink{{Kafka: &KafkaSink{Kafka: Kafka{Topic: sinkTopic}}}},
 			}},
 		},
 	})
 
 	WaitForPipeline()
-
 	WaitForPod()
 
 	n := 500 * 30
 	go PumpKafkaTopic(topic, n)
 
+	stopPortForward := StartPortForward()
+	WaitForNSunkMessages(1000)
+	stopPortForward()
+
 	RestartStatefulSet("kafka-broker")
 	WaitForPod("kafka-broker-0")
 
 	ExpectKafkaTopicCount(sinkTopic, n, 2*time.Minute)
-	defer StartPortForward("kafka-main-0")()
+	defer StartPortForward()()
 	WaitForNoErrors()
 	ExpectLogLine("main", "Failed to connect to broker kafka-broker:9092")
 }
@@ -91,10 +96,8 @@ func TestKafkaFMEA_PipelineDeletedDisruption(t *testing.T) {
 	topic := CreateKafkaTopic()
 	sinkTopic := CreateKafkaTopic()
 
-	time.Sleep(10*time.Second)
-
 	name := CreatePipeline(Pipeline{
-		ObjectMeta: metav1.ObjectMeta{GenerateName: "kafka-"},
+		ObjectMeta: metav1.ObjectMeta{GenerateName: "kafka-pipeline-deleted-disruption-"},
 		Spec: PipelineSpec{
 			Steps: []StepSpec{{
 				Name:    "main",
@@ -111,8 +114,9 @@ func TestKafkaFMEA_PipelineDeletedDisruption(t *testing.T) {
 	n := 500 * 15
 	go PumpKafkaTopic(topic, n)
 
-	defer StartPortForward(name + "-main-0")()
-	WaitForSunkMessages()
+	stopPortForward := StartPortForward()
+	WaitForNSunkMessages(1000)
+	stopPortForward()
 
 	DeletePipelines()
 	WaitForPodsToBeDeleted()
@@ -122,7 +126,7 @@ func TestKafkaFMEA_PipelineDeletedDisruption(t *testing.T) {
 			Steps: []StepSpec{{
 				Name:    "main",
 				Cat:     &Cat{},
-				Sources: []Source{{Kafka: &KafkaSource{Kafka: Kafka{Topic: topic}}}},
+				Sources: []Source{{Kafka: &KafkaSource{StartOffset: "First", Kafka: Kafka{Topic: topic}}}},
 				Sinks:   []Sink{{Kafka: &KafkaSink{Kafka: Kafka{Topic: sinkTopic}}}},
 			}},
 		},
diff --git a/test/kafka-stress/kafka_stress_test.go b/test/kafka-stress/kafka_stress_test.go
index 451f1c03..f6af1552 100644
--- a/test/kafka-stress/kafka_stress_test.go
+++ b/test/kafka-stress/kafka_stress_test.go
@@ -44,7 +44,7 @@ func TestKafkaSourceStress(t *testing.T) {
 
 	WaitForPipeline()
 
-	defer StartPortForward(name + "-main-0")()
+	defer StartPortForward()()
 
 	WaitForPod()
 
@@ -53,8 +53,9 @@ func TestKafkaSourceStress(t *testing.T) {
 
 	defer StartTPSReporter(t, "main", prefix, n)()
 	go PumpKafkaTopic(topic, n, prefix, Params.MessageSize)
-	WaitForPending()
-	WaitForNothingPending()
+	WaitForPending(Params.Timeout)
+	WaitForNothingPending(Params.Timeout)
+	WaitForTotalSourceMessages(n, Params.Timeout)
 	WaitForTotalSunkMessages(n, Params.Timeout)
 }
 
@@ -97,7 +98,7 @@ func testKafkaSinkStress(t *testing.T, async bool) {
 
 	WaitForPipeline()
 
-	defer StartPortForward(name + "-main-0")()
+	defer StartPortForward()()
 
 	WaitForPod()
 
@@ -107,6 +108,8 @@ func testKafkaSinkStress(t *testing.T, async bool) {
 	defer StartTPSReporter(t, "main", prefix, n)()
 
 	go PumpKafkaTopic(topic, n, prefix, Params.MessageSize)
-	WaitForPending()
+	WaitForPending(Params.Timeout)
+	WaitForNothingPending(Params.Timeout)
+	WaitForTotalSourceMessages(n, Params.Timeout)
 	WaitForTotalSunkMessages(n, Params.Timeout)
 }
diff --git a/test/kafka-stress/test-results.json b/test/kafka-stress/test-results.json
index 2deb6a59..9d102b50 100644
--- a/test/kafka-stress/test-results.json
+++ b/test/kafka-stress/test-results.json
@@ -1,13 +1,13 @@
 {
   "TestKafkaAsyncSinkStress/.tps": 400,
-  "TestKafkaSinkStress/.tps": 150,
+  "TestKafkaSinkStress/.tps": 300,
   "TestKafkaSinkStress/N=10,messageSize=100.tps": 200,
   "TestKafkaSinkStress/N=10,messageSize=1000.tps": 150,
   "TestKafkaSinkStress/N=50000.tps": 750,
   "TestKafkaSinkStress/async=true.tps": 400,
   "TestKafkaSinkStress/messageSize=1000.tps": 300,
   "TestKafkaSinkStress/replicas=2.tps": 400,
-  "TestKafkaSourceStress/.tps": 350,
+  "TestKafkaSourceStress/.tps": 400,
   "TestKafkaSourceStress/N=10,messageSize=100.tps": 450,
   "TestKafkaSourceStress/N=10,messageSize=1000.tps": 650,
   "TestKafkaSourceStress/N=50000.tps": 3150,
diff --git a/test/kafka.go b/test/kafka.go
index 6a81f9c4..ad41729d 100644
--- a/test/kafka.go
+++ b/test/kafka.go
@@ -39,7 +39,7 @@ func PumpKafkaTopic(topic string, n int, opts ...interface{}) {
 }
 
 func ExpectKafkaTopicCount(topic string, total int, timeout time.Duration) {
-	var lastCount, stalls int
+	stall := &stall{}
 	log.Printf("expecting %d messages to be sunk to topic %s within %v\n", total, topic, timeout)
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
@@ -57,13 +57,7 @@ func ExpectKafkaTopicCount(topic string, total int, timeout time.Duration) {
 			if count > total {
 				panic(fmt.Errorf("too many messages %d > %d", count, total))
 			}
-			if count == lastCount {
-				stalls++
-			}
-			if stalls >= 3 {
-				panic(fmt.Errorf("stalled waiting for %d messages in topic %q", total, topic))
-			}
-			lastCount = count
+			stall.accept(count)
 			time.Sleep(time.Second)
 		}
 	}
diff --git a/test/log.go b/test/log.go
index d3d3112d..d270837e 100644
--- a/test/log.go
+++ b/test/log.go
@@ -8,10 +8,8 @@ import (
 	"fmt"
 	"log"
 	"regexp"
-	"sync"
 	"time"
 
-	. "github.com/argoproj-labs/argo-dataflow/api/v1alpha1"
 	"github.com/argoproj-labs/argo-dataflow/shared/util"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -121,36 +119,3 @@ func podLogContains(ctx context.Context, podName, containerName string, match fu
 		}
 	}
 }
-
-func TailLogs() {
-	ctx := context.Background()
-	list, err := podsInterface.List(ctx, metav1.ListOptions{LabelSelector: KeyPipelineName})
-	if err != nil {
-		panic(err)
-	}
-	wg := sync.WaitGroup{}
-	for _, item := range list.Items {
-		wg.Add(1)
-		go func(podName string) {
-			defer wg.Done()
-			tailLogs(podName, "init")
-			tailLogs(podName, "sidecar")
-			tailLogs(podName, "main")
-		}(item.Name)
-	}
-	wg.Wait()
-}
-
-func tailLogs(podName, containerName string) {
-	ctx := context.Background()
-	log.Printf("📄 tailing logs for %q/%q\n", podName, containerName)
-	stream, err := podsInterface.GetLogs(podName, &corev1.PodLogOptions{Container: containerName}).Stream(ctx)
-	if err != nil {
-		log.Printf("ERROR: %v\n", err)
-		return
-	}
-	defer func() { _ = stream.Close() }()
-	for s := bufio.NewScanner(stream); s.Scan(); {
-		log.Println("> " + s.Text())
-	}
-}
diff --git a/test/matchers.go b/test/matchers.go
index df00f7bb..79cde61b 100644
--- a/test/matchers.go
+++ b/test/matchers.go
@@ -6,43 +6,43 @@ import "fmt"
 
 type matcher struct {
 	string
-	match func(w float64) bool
+	stall *stall
+	test  func(w int) bool
 }
 
-func (m matcher) String() string { return m.string }
+func (m *matcher) String() string { return m.string }
 
-func Eq(v float64) matcher {
-	return matcher{
+func (m *matcher) match(v int) bool {
+	m.stall.accept(v)
+	return m.test(v)
+}
+
+func Eq(v int) *matcher {
+	return &matcher{
 		fmt.Sprintf("eq %v", v),
-		func(w float64) bool {
+		&stall{},
+		func(w int) bool {
 			return w == v
 		},
 	}
 }
 
-func Missing() matcher {
-	return matcher{
+func Missing() *matcher {
+	return &matcher{
 		"missing",
-		func(w float64) bool {
+		&stall{},
+		func(w int) bool {
 			return w == missing
 		},
 	}
 }
 
-func Gt(v float64) matcher {
-	return matcher{
+func Gt(v int) *matcher {
+	return &matcher{
 		fmt.Sprintf("gt %v", v),
-		func(w float64) bool {
+		&stall{},
+		func(w int) bool {
 			return w > v
 		},
 	}
 }
-
-func Between(min, max float64) matcher {
-	return matcher{
-		fmt.Sprintf("between %v and <=%v", min, max),
-		func(w float64) bool {
-			return min <= w && w <= max
-		},
-	}
-}
diff --git a/test/metric.go b/test/metric.go
index 25594358..d6677b41 100644
--- a/test/metric.go
+++ b/test/metric.go
@@ -14,16 +14,16 @@ import (
 	"github.com/prometheus/common/expfmt"
 )
 
-func WaitForPending() {
-	ExpectMetric("sources_pending", Gt(0))
+func WaitForPending(opts ...interface{}) {
+	ExpectMetric("sources_pending", Gt(0), opts...)
 }
 
-func WaitForNothingPending() {
-	ExpectMetric("sources_pending", Eq(0))
+func WaitForNothingPending(opts ...interface{}) {
+	ExpectMetric("sources_pending", Eq(0), opts...)
 }
 
-func WaitForTotalSourceMessages(v int) {
-	ExpectMetric("sources_total", Eq(float64(v)))
+func WaitForTotalSourceMessages(v int, opts ...interface{}) {
+	ExpectMetric("sources_total", Eq(v), opts...)
 }
 
 func WaitForNoErrors() {
@@ -31,20 +31,24 @@ func WaitForNoErrors() {
 }
 
 func WaitForSunkMessages(opts ...interface{}) {
-	ExpectMetric("sinks_total", Gt(0), opts...)
+	WaitForNSunkMessages(0, opts...)
+}
+
+func WaitForNSunkMessages(v int, opts ...interface{}) {
+	ExpectMetric("sinks_total", Gt(v), opts...)
 }
 
 func WaitForTotalSunkMessages(v int, opts ...interface{}) {
-	ExpectMetric("sinks_total", Eq(float64(v)), opts...)
+	ExpectMetric("sinks_total", Eq(v), opts...)
 }
 
 func WaitForProcessLatencySeconds(v int, opts ...interface{}) {
-	ExpectMetric("process_latency_seconds", Eq(float64(v)), opts...)
+	ExpectMetric("process_latency_seconds", Eq(v), opts...)
 }
 
-var missing = rand.Float64()
+var missing = rand.Int()
 
-func ExpectMetric(name string, matcher matcher, opts ...interface{}) {
+func ExpectMetric(name string, matcher *matcher, opts ...interface{}) {
 	ctx := context.Background()
 	port := 3569
 	timeout := 30 * time.Second
@@ -72,7 +76,6 @@ func ExpectMetric(name string, matcher matcher, opts ...interface{}) {
 					found = true
 					for _, m := range family.Metric {
 						v := getValue(m)
-						fmt.Println(v)
 						if matcher.match(v) {
 							return
 						} else {
@@ -89,11 +92,11 @@ func ExpectMetric(name string, matcher matcher, opts ...interface{}) {
 	}
 }
 
-func getValue(m *io_prometheus_client.Metric) float64 {
+func getValue(m *io_prometheus_client.Metric) int {
 	if x := m.Counter; x != nil {
-		return x.GetValue()
+		return int(x.GetValue())
 	} else if x := m.Gauge; x != nil {
-		return x.GetValue()
+		return int(x.GetValue())
 	} else {
 		panic(fmt.Errorf("metric not-supported (not a counter/gauge)"))
 	}
diff --git a/test/pod.go b/test/pod.go
index a44b5858..ca8dff1c 100644
--- a/test/pod.go
+++ b/test/pod.go
@@ -55,7 +55,7 @@ func WaitForPodsToBeDeleted() {
 	}
 }
 
-func WaitForPod(opts ...interface{}) {
+func WaitForPod(opts ...interface{}) (podName string) {
 	// by default, wait for any pod to be ready
 	var (
 		listOptions = metav1.ListOptions{LabelSelector: KeyPipelineName}
@@ -101,7 +101,7 @@ func WaitForPod(opts ...interface{}) {
 			}
 			log.Printf("pod %q is %s %q %q\n", p.Name, s.Phase, s.Message, y)
 			if f(p) {
-				return
+				return p.Name
 			}
 		}
 	}
diff --git a/test/port_forward.go b/test/port_forward.go
index dc5b2c94..6b1b89d0 100644
--- a/test/port_forward.go
+++ b/test/port_forward.go
@@ -14,18 +14,24 @@ import (
 	"k8s.io/client-go/transport/spdy"
 )
 
-func StartPortForward(podName string, opts ...interface{}) (stopPortForward func()) {
-	WaitForPod(podName)
-
+func StartPortForward(opts ...interface{}) (stopPortForward func()) {
 	port := 3569
+	podName := ""
 	for _, opt := range opts {
 		switch v := opt.(type) {
 		case int:
 			port = v
+		case string:
+			podName = v
 		default:
 			panic("unknown option")
 		}
 	}
+	if podName != "" {
+		WaitForPod(podName)
+	} else {
+		podName = WaitForPod()
+	}
 	log.Printf("starting port-forward to pod %q on %d\n", podName, port)
 	transport, upgrader, err := spdy.RoundTripperFor(restConfig)
 	if err != nil {
diff --git a/test/stall.go b/test/stall.go
new file mode 100644
index 00000000..3b1e0084
--- /dev/null
+++ b/test/stall.go
@@ -0,0 +1,20 @@
+package test
+
+import "fmt"
+
+type stall struct {
+	stalls int
+	last   int
+}
+
+func (s *stall) accept(v int) {
+	if v == s.last {
+		s.stalls++
+	} else {
+		s.stalls = 0
+	}
+	s.last = v
+	if s.stalls >= 10 {
+		panic(fmt.Errorf("stalled at %d", v))
+	}
+}
diff --git a/testapi/kafka.go b/testapi/kafka.go
index b79d5364..e08d71ca 100644
--- a/testapi/kafka.go
+++ b/testapi/kafka.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"strconv"
@@ -11,15 +12,13 @@ import (
 )
 
 func init() {
-	config := &kafka.ConfigMap{
-		"bootstrap.servers": "kafka-broker:9092",
-		"group.id":          "testapi",
-	}
-
+	const bootstrapServers = "kafka-broker:9092"
 	http.HandleFunc("/kafka/create-topic", func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
 		topic := r.URL.Query().Get("topic")
-		admin, err := kafka.NewAdminClient(config)
+		admin, err := kafka.NewAdminClient(&kafka.ConfigMap{
+			"bootstrap.servers": bootstrapServers,
+		})
 		if err != nil {
 			w.WriteHeader(500)
 			_, _ = w.Write([]byte(err.Error()))
@@ -38,39 +37,56 @@ func init() {
 		w.WriteHeader(201)
 	})
 	http.HandleFunc("/kafka/count-topic", func(w http.ResponseWriter, r *http.Request) {
-		topics := r.URL.Query()["topic"]
-		if len(topics) < 1 {
-			w.WriteHeader(400)
-			return
-		}
-		topic := topics[0]
-		consumer, err := kafka.NewConsumer(config)
+		topic := r.URL.Query().Get("topic")
+		consumer, err := kafka.NewConsumer(&kafka.ConfigMap{
+			"bootstrap.servers":        bootstrapServers,
+			"group.id":                 "testapi",
+			"auto.offset.reset":        "earliest",
+			"enable.auto.commit":       false,
+			"enable.auto.offset.store": false,
+			"statistics.interval.ms":   500,
+		})
 		if err != nil {
 			w.WriteHeader(500)
 			_, _ = w.Write([]byte(err.Error()))
 			return
 		}
 		defer consumer.Close()
-		md, err := consumer.GetMetadata(&topic, false, 3000)
+		err = consumer.Subscribe(topic, nil)
 		if err != nil {
 			w.WriteHeader(500)
 			_, _ = w.Write([]byte(err.Error()))
 			return
 		}
-		count := 0
-		for _, t := range md.Topics {
-			for _, p := range t.Partitions {
-				_, h, err := consumer.QueryWatermarkOffsets(topic, p.ID, 3000)
-				if err != nil {
+		for {
+			ev := consumer.Poll(5 * 1000)
+			select {
+			case <-r.Context().Done():
+				return
+			default:
+				switch e := ev.(type) {
+				case *kafka.Message:
+				case *kafka.Stats:
+					// https://github.com/edenhill/librdkafka/wiki/Consumer-lag-monitoring
+					// https://github.com/confluentinc/confluent-kafka-go/blob/master/examples/stats_example/stats_example.go
+					println(ev.String())
+					stats := &KafkaStats{}
+					if err := json.Unmarshal([]byte(e.String()), stats); err != nil {
+						w.WriteHeader(500)
+						_, _ = fmt.Fprintf(w, "failed to unmarshall stats: %v\n", err)
+					} else {
+						w.WriteHeader(200)
+						_, _ = w.Write([]byte(fmt.Sprint(stats.count(topic))))
+					}
+					return
+				case nil:
+				default:
 					w.WriteHeader(500)
-					_, _ = w.Write([]byte(err.Error()))
+					_, _ = w.Write([]byte(fmt.Sprint(e)))
 					return
 				}
-				count += int(h)
 			}
 		}
-		w.WriteHeader(200)
-		_, _ = w.Write([]byte(strconv.Itoa(count)))
 	})
 	http.HandleFunc("/kafka/pump-topic", func(w http.ResponseWriter, r *http.Request) {
 		topic := r.URL.Query().Get("topic")
@@ -96,7 +112,9 @@ func init() {
 		w.Header().Set("Content-Type", "application/octet-stream")
 		w.WriteHeader(200)
 
-		producer, err := kafka.NewProducer(config)
+		producer, err := kafka.NewProducer(&kafka.ConfigMap{
+			"bootstrap.servers": bootstrapServers,
+		})
 		if err != nil {
 			w.WriteHeader(500)
 			_, _ = w.Write([]byte(err.Error()))
diff --git a/testapi/kafka_stats.go b/testapi/kafka_stats.go
new file mode 100644
index 00000000..50bdb7f6
--- /dev/null
+++ b/testapi/kafka_stats.go
@@ -0,0 +1,18 @@
+package main
+
+type KafkaStats struct {
+	Topics map[string]struct {
+		Partitions map[string]struct {
+			LoOffset int64 `json:"lo_offset"`
+			HiOffset int64 `json:"hi_offset"`
+		} `json:"partitions"`
+	} `json:"topics"`
+}
+
+func (s KafkaStats) count(topic string) int64 {
+	var count int64
+	for _, p := range s.Topics[topic].Partitions {
+		count += p.HiOffset - p.LoOffset
+	}
+	return count
+}