fix(output): drain buffered records on TCP/UDP Stop (PIPE-1230)

output/tcp/tcp.go

@@ -130,16 +130,18 @@ func (t *TCP) Stop(ctx context.Context) error {
 	// Record zero active workers
 	output.BlitzOutputActiveWorkersGauge.Record(ctx, 0, outputType)
 
-	// Close the channel to ensure workers do not
-	// process new data.
-	close(t.dataChan)
-
-	// Signal the workers to stop.
+	// Reject any further writes and stop the workers (and their restart loop).
 	t.cancel()
-
-	// Stop the worker manager
 	t.workerManager.Stop()
 
+	// The workers have now exited, so this goroutine is the sole owner of
+	// dataChan. Close it and drain any records that were still buffered when
+	// shutdown began, delivering them instead of dropping them (PIPE-1230).
+	// Draining here, after the workers are joined, keeps delivery deterministic
+	// rather than racing a worker's channel-drain against context cancellation.
+	close(t.dataChan)
+	t.drainBuffered(ctx)
+
 	t.logger.Info("TCP output stopped successfully")
 	return nil
 }
@@ -183,6 +185,48 @@ func (t *TCP) tcpWorker(id int) {
 	}
 }
 
+// drainBuffered sends any records still buffered in dataChan at shutdown. It runs
+// after the workers have stopped, so it is the sole consumer of the (now closed)
+// channel and delivery is deterministic. It is bounded by the connect/write
+// timeouts, DefaultTCPStopTimeout, and the caller's context so an unreachable
+// destination cannot hang shutdown.
+func (t *TCP) drainBuffered(ctx context.Context) {
+	if len(t.dataChan) == 0 {
+		return
+	}
+
+	conn, err := t.connect()
+	if err != nil {
+		t.logger.Error("Failed to connect while draining buffered records on shutdown",
+			zap.Int("dropped", len(t.dataChan)),
+			zap.Error(err))
+		return
+	}
+	defer conn.Close()
+
+	t.drainTo(ctx, conn, time.Now().Add(DefaultTCPStopTimeout))
+}
+
+// drainTo sends every record remaining in dataChan over conn until the channel
+// is closed and empty, a send fails, or ctx is done / the deadline passes. It is
+// split out from drainBuffered so the send-failure and deadline paths can be
+// exercised directly with a fake connection.
+func (t *TCP) drainTo(ctx context.Context, conn net.Conn, deadline time.Time) {
+	for data := range t.dataChan {
+		if ctx.Err() != nil || time.Now().After(deadline) {
+			t.logger.Warn("Shutdown deadline reached while draining buffered records",
+				zap.Int("dropped", len(t.dataChan)+1))
+			return
+		}
+		if err := t.sendData(conn, data); err != nil {
+			t.logger.Error("Failed to send buffered record while draining on shutdown",
+				zap.Int("dropped", len(t.dataChan)+1),
+				zap.Error(err))
+			return
+		}
+	}
+}
+
 // connect establishes a TCP connection to the configured host and port
 func (t *TCP) connect() (net.Conn, error) {
 	address := net.JoinHostPort(t.host, t.port)

output/tcp/tcp_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"fmt"
+	"io"
 	"net"
 	"os"
 	"path/filepath"
@@ -473,6 +475,40 @@ func TestTCP_IntegrationTLS(t *testing.T) {
 	}
 }
 
+// TestTCP_StopWithUnreachableDestinationDoesNotHang verifies that Stop completes
+// (without hanging) when buffered records cannot be drained because the
+// destination is unreachable. Pointing at a closed port means the worker never
+// connects, so the records stay buffered until Stop, and the drain's connect
+// attempt fails fast rather than blocking shutdown.
+func TestTCP_StopWithUnreachableDestinationDoesNotHang(t *testing.T) {
+	logger := zap.NewNop()
+
+	// Reserve a port, then close the listener so nothing is listening on it.
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	_, port, err := net.SplitHostPort(l.Addr().String())
+	require.NoError(t, err)
+	require.NoError(t, l.Close())
+
+	tcp, err := New(logger, "127.0.0.1", port, 1, nil)
+	require.NoError(t, err)
+
+	// The worker cannot connect, so these stay buffered (well under channel cap).
+	ctx := context.Background()
+	for i := 0; i < 10; i++ {
+		require.NoError(t, tcp.Write(ctx, output.LogRecord{Message: fmt.Sprintf("unreachable-%d", i)}))
+	}
+
+	done := make(chan error, 1)
+	go func() { done <- tcp.Stop(ctx) }()
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(15 * time.Second):
+		t.Fatal("Stop hung when draining to an unreachable destination")
+	}
+}
+
 // Test server implementation
 var (
 	receivedData [][]byte
@@ -540,3 +576,108 @@ func getReceivedData(t *testing.T) [][]byte {
 
 	return result
 }
+
+// fakeConn is a minimal net.Conn used to exercise drainTo's send-failure and
+// context/deadline branches without real network I/O.
+type fakeConn struct {
+	writeErr error
+	writes   int
+}
+
+func (f *fakeConn) Read([]byte) (int, error) { return 0, io.EOF }
+func (f *fakeConn) Write(b []byte) (int, error) {
+	f.writes++
+	if f.writeErr != nil {
+		return 0, f.writeErr
+	}
+	return len(b), nil
+}
+func (f *fakeConn) Close() error                     { return nil }
+func (f *fakeConn) LocalAddr() net.Addr              { return &net.TCPAddr{} }
+func (f *fakeConn) RemoteAddr() net.Addr             { return &net.TCPAddr{} }
+func (f *fakeConn) SetDeadline(time.Time) error      { return nil }
+func (f *fakeConn) SetReadDeadline(time.Time) error  { return nil }
+func (f *fakeConn) SetWriteDeadline(time.Time) error { return nil }
+
+// drainBuffered with an empty channel must return before attempting to connect.
+func TestTCP_drainBuffered_emptyChannelReturnsEarly(t *testing.T) {
+	tcp := &TCP{logger: zap.NewNop(), dataChan: make(chan string, 1)}
+	tcp.drainBuffered(context.Background())
+}
+
+// drainTo stops at the first send failure, leaving the rest undrained.
+func TestTCP_drainTo_stopsOnSendError(t *testing.T) {
+	tcp := &TCP{logger: zap.NewNop(), dataChan: make(chan string, 4)}
+	tcp.dataChan <- "one"
+	tcp.dataChan <- "two"
+	close(tcp.dataChan)
+
+	conn := &fakeConn{writeErr: fmt.Errorf("write failed")}
+	tcp.drainTo(context.Background(), conn, time.Now().Add(time.Hour))
+
+	require.Equal(t, 1, conn.writes, "drainTo should stop after the first failed send")
+}
+
+// drainTo returns without sending when the context is already done.
+func TestTCP_drainTo_stopsWhenContextDone(t *testing.T) {
+	tcp := &TCP{logger: zap.NewNop(), dataChan: make(chan string, 4)}
+	tcp.dataChan <- "one"
+	close(tcp.dataChan)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	conn := &fakeConn{}
+	tcp.drainTo(ctx, conn, time.Now().Add(time.Hour))
+
+	require.Equal(t, 0, conn.writes, "drainTo should not send once the context is done")
+}
+
+// TestTCP_StopDrainsBufferedRecords verifies Stop() drains everything already
+// buffered before returning, instead of dropping records that are still queued
+// when shutdown begins. Unlike the other delivery tests, this one does NOT poll
+// for delivery before calling Stop — draining is Stop()'s responsibility.
+func TestTCP_StopDrainsBufferedRecords(t *testing.T) {
+	logger := zap.NewNop()
+
+	listener, serverAddr := startTestTCPServer(t)
+	defer listener.Close()
+
+	host, port, err := net.SplitHostPort(serverAddr)
+	require.NoError(t, err)
+
+	tcp, err := New(logger, host, port, 1, nil)
+	require.NoError(t, err)
+
+	// Queue a backlog of uniquely-identifiable records as fast as possible, so a
+	// large number are still buffered when Stop is called, then stop immediately.
+	const n = 100
+	ctx := context.Background()
+	for i := 0; i < n; i++ {
+		require.NoError(t, tcp.Write(ctx, output.LogRecord{Message: fmt.Sprintf("drain-msg-%d", i)}))
+	}
+
+	require.NoError(t, tcp.Stop(ctx))
+
+	// Once Stop has returned, every record must have been delivered. Poll only to
+	// let the test server goroutine finish reading bytes already on the wire; no
+	// new records can be sent after Stop returns, so a missing record means it was
+	// dropped rather than drained.
+	// Once Stop has returned, every record must have been delivered. Poll only to
+	// let the test server goroutine finish reading bytes already on the wire; no
+	// new records can be sent after Stop returns, so a missing record means it was
+	// dropped rather than drained.
+	require.Eventually(t, func() bool {
+		var all []byte
+		for _, d := range getReceivedData(t) {
+			all = append(all, d...)
+		}
+		s := string(all)
+		for i := 0; i < n; i++ {
+			if !strings.Contains(s, fmt.Sprintf("drain-msg-%d\n", i)) {
+				return false
+			}
+		}
+		return true
+	}, 3*time.Second, 10*time.Millisecond, "all buffered records should be delivered before Stop returns")
+}

output/udp/udp.go

@@ -123,20 +123,64 @@ func (u *UDP) Stop(ctx context.Context) error {
 	// Record zero active workers
 	output.BlitzOutputActiveWorkersGauge.Record(ctx, 0, outputType)
 
-	// Close the channel to ensure workers do not
-	// process new data.
-	close(u.dataChan)
-
-	// Signal the workers to stop.
+	// Reject any further writes and stop the workers (and their restart loop).
 	u.cancel()
-
-	// Stop the worker manager
 	u.workerManager.Stop()
 
+	// The workers have now exited, so this goroutine is the sole owner of
+	// dataChan. Close it and drain any records that were still buffered when
+	// shutdown began, delivering them instead of dropping them (PIPE-1230).
+	// Draining here, after the workers are joined, keeps delivery deterministic
+	// rather than racing a worker's channel-drain against context cancellation.
+	close(u.dataChan)
+	u.drainBuffered(ctx)
+
 	u.logger.Info("UDP output stopped successfully")
 	return nil
 }
 
+// drainBuffered sends any records still buffered in dataChan at shutdown. It runs
+// after the workers have stopped, so it is the sole consumer of the (now closed)
+// channel and delivery is deterministic. It is bounded by the connect/write
+// timeouts, DefaultUDPStopTimeout, and the caller's context so an unreachable
+// destination cannot hang shutdown.
+func (u *UDP) drainBuffered(ctx context.Context) {
+	if len(u.dataChan) == 0 {
+		return
+	}
+
+	conn, err := u.connect()
+	if err != nil {
+		u.logger.Error("Failed to connect while draining buffered records on shutdown",
+			zap.Int("dropped", len(u.dataChan)),
+			zap.Error(err))
+		return
+	}
+	defer conn.Close()
+
+	u.drainTo(ctx, conn, time.Now().Add(DefaultUDPStopTimeout))
+}
+
+// drainTo sends every record remaining in dataChan over conn until the channel
+// is closed and empty, a send fails, or ctx is done / the deadline passes. It is
+// split out from drainBuffered so the send-failure and deadline paths can be
+// exercised directly with a fake connection.
+func (u *UDP) drainTo(ctx context.Context, conn net.Conn, deadline time.Time) {
+	for data := range u.dataChan {
+		if ctx.Err() != nil || time.Now().After(deadline) {
+			u.logger.Warn("Shutdown deadline reached while draining buffered records",
+				zap.Int("dropped", len(u.dataChan)+1))
+			return
+		}
+		if err := u.sendData(conn, data); err != nil {
+			u.logger.Error("Failed to send buffered record while draining on shutdown",
+				zap.Int("dropped", len(u.dataChan)+1),
+				zap.Error(err))
+			return
+		}
+	}
+}
+
 // udpWorker processes UDP data from the channel and sends it to the configured host and port.
 // This function is designed to work with the worker manager, which handles automatic restart
 // with exponential backoff when the worker exits due to connection failures or errors.