Update probe

Author: ilyario

cmd/processor/main.go

@@ -146,19 +146,7 @@ func main() {
 	})
 	progressTimeout := progressTimeoutFromEnv(os.Getenv("PROCESSOR_PROGRESS_TIMEOUT_SECONDS"))
 	mux.HandleFunc("/livez", func(w http.ResponseWriter, _ *http.Request) {
-		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
-		if !proc.Ready() {
-			w.WriteHeader(http.StatusServiceUnavailable)
-			_, _ = w.Write([]byte("not ready\n"))
-			return
-		}
-		if proc.ProgressStale(progressTimeout) {
-			w.WriteHeader(http.StatusServiceUnavailable)
-			_, _ = w.Write([]byte("stale: no pipeline progress\n"))
-			return
-		}
-		w.WriteHeader(http.StatusOK)
-		_, _ = w.Write([]byte("ok\n"))
+		writeLivez(w, proc.Ready, proc.ProgressStale, progressTimeout)
 	})
 	metricsServer := &http.Server{Addr: metricsPort, Handler: mux}
 	go func() {
@@ -230,6 +218,23 @@ func main() {
 	logger.Info("Processor stopped successfully")
 }
 
+// writeLivez implements GET /livez for Kubernetes liveness probes.
+func writeLivez(w http.ResponseWriter, ready func() bool, progressStale func(time.Duration) bool, progressTimeout time.Duration) {
+	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+	if !ready() {
+		w.WriteHeader(http.StatusServiceUnavailable)
+		_, _ = w.Write([]byte("not ready\n"))
+		return
+	}
+	if progressStale(progressTimeout) {
+		w.WriteHeader(http.StatusServiceUnavailable)
+		_, _ = w.Write([]byte("stale: no pipeline progress\n"))
+		return
+	}
+	w.WriteHeader(http.StatusOK)
+	_, _ = w.Write([]byte("ok\n"))
+}
+
 // processorLevelFromEnv returns zap LevelEnabler from LOG_LEVEL env if set, otherwise optsLevel.
 // progressTimeoutFromEnv parses PROCESSOR_PROGRESS_TIMEOUT_SECONDS (0 or unset disables stale checks).
 func progressTimeoutFromEnv(env string) time.Duration {

cmd/processor/main_test.go

@@ -3,7 +3,10 @@ package main
 import (
 	"context"
 	"errors"
+	"net/http"
+	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -14,6 +17,41 @@ func TestIsRetryableProcessorError(t *testing.T) {
 	assert.False(t, isRetryableProcessorError(errors.New("permission denied")))
 }
 
+func TestWriteLivez(t *testing.T) {
+	t.Parallel()
+
+	t.Run("not ready", func(t *testing.T) {
+		t.Parallel()
+		rr := httptest.NewRecorder()
+		writeLivez(rr, func() bool { return false }, func(time.Duration) bool { return false }, time.Minute)
+		assert.Equal(t, http.StatusServiceUnavailable, rr.Code)
+		assert.Contains(t, rr.Body.String(), "not ready")
+	})
+
+	t.Run("stale progress", func(t *testing.T) {
+		t.Parallel()
+		rr := httptest.NewRecorder()
+		writeLivez(rr, func() bool { return true }, func(time.Duration) bool { return true }, time.Minute)
+		assert.Equal(t, http.StatusServiceUnavailable, rr.Code)
+		assert.Contains(t, rr.Body.String(), "stale")
+	})
+
+	t.Run("ok", func(t *testing.T) {
+		t.Parallel()
+		rr := httptest.NewRecorder()
+		writeLivez(rr, func() bool { return true }, func(time.Duration) bool { return false }, time.Minute)
+		assert.Equal(t, http.StatusOK, rr.Code)
+		assert.Equal(t, "ok\n", rr.Body.String())
+	})
+}
+
+func TestProgressTimeoutFromEnv(t *testing.T) {
+	t.Parallel()
+	assert.Equal(t, 10*time.Minute, progressTimeoutFromEnv(""))
+	assert.Equal(t, time.Duration(0), progressTimeoutFromEnv("0"))
+	assert.Equal(t, 30*time.Minute, progressTimeoutFromEnv("1800"))
+}
+
 func TestProcessorSinkErrorMaxRetriesFromEnv(t *testing.T) {
 	assert.Equal(t, 0, processorSinkErrorMaxRetriesFromEnv(""))
 	assert.Equal(t, 0, processorSinkErrorMaxRetriesFromEnv("-1"))

internal/connectors/base.go

@@ -185,6 +185,21 @@ func (p *progressRecorder) notifyProgress() {
 	}
 }
 
+// AckMessages calls Ack on each message (commits source offsets when configured).
+func AckMessages(msgs []*types.Message) {
+	for _, m := range msgs {
+		if m.Ack != nil {
+			m.Ack()
+		}
+	}
+}
+
+// AckMessagesAndNotifyProgress commits offsets and updates liveness progress after a successful batch.
+func (p *progressRecorder) AckMessagesAndNotifyProgress(msgs []*types.Message) {
+	AckMessages(msgs)
+	p.notifyProgress()
+}
+
 // SetConnectorInfo sets the connector type and role for metrics.
 func (c *connectorMetadata) SetConnectorInfo(connectorType, role string) {
 	c.connectorType = connectorType

internal/connectors/base_progress_test.go

@@ -0,0 +1,58 @@
+/*
+Copyright 2024.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package connectors
+
+import (
+	"sync/atomic"
+	"testing"
+
+	"github.com/dataflow-operator/dataflow/internal/types"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAckMessagesAndNotifyProgress(t *testing.T) {
+	t.Parallel()
+
+	var progressCalls atomic.Int32
+	var ackCalls atomic.Int32
+
+	rec := &progressRecorder{}
+	rec.SetProgressCallback(func() {
+		progressCalls.Add(1)
+	})
+
+	msg := types.NewMessage([]byte(`{}`))
+	msg.Ack = func() { ackCalls.Add(1) }
+
+	rec.AckMessagesAndNotifyProgress([]*types.Message{msg})
+
+	assert.Equal(t, int32(1), ackCalls.Load())
+	assert.Equal(t, int32(1), progressCalls.Load())
+}
+
+func TestAckMessages_NoAckCallback(t *testing.T) {
+	t.Parallel()
+
+	var progressCalls atomic.Int32
+	rec := &progressRecorder{}
+	rec.SetProgressCallback(func() {
+		progressCalls.Add(1)
+	})
+
+	rec.AckMessagesAndNotifyProgress([]*types.Message{types.NewMessage([]byte(`{}`))})
+	assert.Equal(t, int32(1), progressCalls.Load())
+}

internal/connectors/batch_writer_test.go

@@ -174,6 +174,37 @@ func TestRunBatchWriteLoop_InFlightCancel(t *testing.T) {
 	assert.True(t, gotDetached.Load())
 }
 
+func TestRunBatchWriteLoop_OnAck_NotifiesProgress(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	ch := make(chan *types.Message, 1)
+	ch <- types.NewMessage([]byte(`{}`))
+	close(ch)
+
+	var progressCalls int
+	sink := &progressRecorder{}
+	sink.SetProgressCallback(func() {
+		progressCalls++
+	})
+
+	err := RunBatchWriteLoop(ctx, ch, BatchWriteConfig{MaxBatchSize: 1, FlushInterval: 0}, BatchWriteOptions{
+		Logger: logr.Discard(),
+		OnFlush: func(_ context.Context, _ []*types.Message) error {
+			return nil
+		},
+		OnAck: func(msgs []*types.Message) {
+			sink.AckMessagesAndNotifyProgress(msgs)
+		},
+	})
+	if err != nil {
+		t.Fatalf("RunBatchWriteLoop: %v", err)
+	}
+	if progressCalls != 1 {
+		t.Fatalf("progressCalls = %d, want 1", progressCalls)
+	}
+}
+
 func TestRunBatchWriteLoop_OnAck(t *testing.T) {
 	ctx := context.Background()
 	ch := make(chan *types.Message, 1)

internal/connectors/clickhouse.go

@@ -254,6 +254,7 @@ type ClickHouseSinkConnector struct {
 	baseConnector
 	connectorLogger
 	connectorMetadata
+	progressRecorder
 	rawModeConfig
 	flattenMetadataSinkState
 	config         *v1.ClickHouseSinkSpec
@@ -734,6 +735,9 @@ func (c *ClickHouseSinkConnector) Write(ctx context.Context, messages <-chan *ty
 			c.logger.V(1).Info("Received message for ClickHouse", logkeys.MessageID, types.MessageID(msg), "messageNumber", messageCount, "table", c.config.Table, "fields", getKeys(data))
 			return true
 		},
+		OnAck: func(msgs []*types.Message) {
+			c.AckMessagesAndNotifyProgress(msgs)
+		},
 	})
 }
 

internal/connectors/kafka.go

@@ -47,10 +47,14 @@ import (
 )
 
 // KafkaSourceConnector implements SourceConnector for Kafka
+// defaultKafkaSourceProgressHeartbeat is how often an active consumer group session updates liveness while idle.
+const defaultKafkaSourceProgressHeartbeat = 60 * time.Second
+
 type KafkaSourceConnector struct {
 	baseConnector
 	connectorLogger
 	connectorMetadata
+	progressRecorder
 	config            *v1.KafkaSourceSpec
 	consumer          sarama.ConsumerGroup
 	channelBufferSize int
@@ -65,6 +69,8 @@ type KafkaSourceConnector struct {
 	testConsumeFunc func(ctx context.Context, topics []string, handler sarama.ConsumerGroupHandler) error
 	// readErrCh receives fatal errors from the consumer after Read returns; set during Read.
 	readErrCh chan error
+	// progressHeartbeatInterval overrides defaultKafkaSourceProgressHeartbeat when > 0 (tests only).
+	progressHeartbeatInterval time.Duration
 }
 
 // schemaCache caches Avro schemas by ID
@@ -792,6 +798,13 @@ type kafkaConsumerGroupHandler struct {
 	readyOnce sync.Once // Protects ready channel from being closed multiple times
 }
 
+func (k *KafkaSourceConnector) kafkaProgressHeartbeatInterval() time.Duration {
+	if k.progressHeartbeatInterval > 0 {
+		return k.progressHeartbeatInterval
+	}
+	return defaultKafkaSourceProgressHeartbeat
+}
+
 func (h *kafkaConsumerGroupHandler) Setup(sarama.ConsumerGroupSession) error {
 	// Use sync.Once to ensure channel is closed only once
 	// This protects against multiple Setup calls during rebalancing
@@ -811,6 +824,7 @@ func (h *kafkaConsumerGroupHandler) Setup(sarama.ConsumerGroupSession) error {
 		}()
 		close(h.ready)
 	})
+	h.connector.notifyProgress()
 	return nil
 }
 
@@ -823,6 +837,9 @@ func (h *kafkaConsumerGroupHandler) Cleanup(sarama.ConsumerGroupSession) error {
 const kafkaMarkChannelBuffer = 256
 
 func (h *kafkaConsumerGroupHandler) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
+	heartbeat := time.NewTicker(h.connector.kafkaProgressHeartbeatInterval())
+	defer heartbeat.Stop()
+
 	markChan := make(chan *sarama.ConsumerMessage, kafkaMarkChannelBuffer)
 
 	markPending := func(message *sarama.ConsumerMessage) {
@@ -845,6 +862,8 @@ func (h *kafkaConsumerGroupHandler) ConsumeClaim(session sarama.ConsumerGroupSes
 
 	for {
 		select {
+		case <-heartbeat.C:
+			h.connector.notifyProgress()
 		case message := <-claim.Messages():
 			if message == nil {
 				drainMarks()
@@ -907,6 +926,7 @@ type KafkaSinkConnector struct {
 	baseConnector
 	connectorLogger
 	connectorMetadata
+	progressRecorder
 	config   *v1.KafkaSinkSpec
 	producer sarama.SyncProducer
 }
@@ -1059,6 +1079,7 @@ func (k *KafkaSinkConnector) Write(ctx context.Context, messages <-chan *types.M
 			if msg.Ack != nil {
 				msg.Ack()
 			}
+			k.notifyProgress()
 		}
 	}
 }