diff --git a/output/tcp/tcp.go b/output/tcp/tcp.go index d36a43d..4068a85 100644 --- a/output/tcp/tcp.go +++ b/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) diff --git a/output/tcp/tcp_test.go b/output/tcp/tcp_test.go index 99c6344..cacf169 100644 --- a/output/tcp/tcp_test.go +++ b/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") +} diff --git a/output/udp/udp.go b/output/udp/udp.go index 1150d49..7592bf7 100644 --- a/output/udp/udp.go +++ b/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. diff --git a/output/udp/udp_test.go b/output/udp/udp_test.go index 035ee62..5ccb52d 100644 --- a/output/udp/udp_test.go +++ b/output/udp/udp_test.go @@ -2,6 +2,8 @@ package udp import ( "context" + "fmt" + "io" "net" "strings" "sync" @@ -13,6 +15,63 @@ import ( "go.uber.org/zap" ) +// 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.UDPAddr{} } +func (f *fakeConn) RemoteAddr() net.Addr { return &net.UDPAddr{} } +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 } + +func TestUDP_drainBuffered_emptyChannelReturnsEarly(t *testing.T) { + u := &UDP{logger: zap.NewNop(), dataChan: make(chan string, 1)} + u.drainBuffered(context.Background()) +} + +func TestUDP_drainBuffered_connectErrorReturns(t *testing.T) { + u := &UDP{logger: zap.NewNop(), host: "nonexistent.invalid", port: "1", dataChan: make(chan string, 1)} + u.dataChan <- "x" + u.drainBuffered(context.Background()) +} + +func TestUDP_drainTo_stopsOnSendError(t *testing.T) { + u := &UDP{logger: zap.NewNop(), dataChan: make(chan string, 2)} + u.dataChan <- "one" + u.dataChan <- "two" + close(u.dataChan) + + conn := &fakeConn{writeErr: fmt.Errorf("write failed")} + u.drainTo(context.Background(), conn, time.Now().Add(time.Hour)) + require.Equal(t, 1, conn.writes) +} + +func TestUDP_drainTo_stopsWhenContextDone(t *testing.T) { + u := &UDP{logger: zap.NewNop(), dataChan: make(chan string, 1)} + u.dataChan <- "one" + close(u.dataChan) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + conn := &fakeConn{} + u.drainTo(ctx, conn, time.Now().Add(time.Hour)) + require.Equal(t, 0, conn.writes) +} + func TestNew(t *testing.T) { logger := zap.NewNop() @@ -304,6 +363,37 @@ func TestUDP_StopTwice(t *testing.T) { udp.Stop(ctx) } +func TestUDP_StopDrainsBufferedRecords(t *testing.T) { + logger := zap.NewNop() + listener, serverAddr := startTestUDPServer(t) + defer listener.Close() + host, port, err := net.SplitHostPort(serverAddr) + require.NoError(t, err) + udp, err := New(logger, host, port, 1) + require.NoError(t, err) + const n = 100 + ctx := context.Background() + for i := 0; i < n; i++ { + require.NoError(t, udp.Write(ctx, output.LogRecord{Message: fmt.Sprintf("drain-msg-%d", i)})) + } + require.NoError(t, udp.Stop(ctx)) + require.Eventually(t, func() bool { + // UDP sends one datagram per record and does not append a newline, so each + // buffered record arrives as its own datagram. Match each expected message + // exactly to disambiguate e.g. drain-msg-1 from drain-msg-10. + received := make(map[string]bool) + for _, d := range getReceivedUDPData(t) { + received[string(d)] = true + } + for i := 0; i < n; i++ { + if !received[fmt.Sprintf("drain-msg-%d", i)] { + return false + } + } + return true + }, 3*time.Second, 10*time.Millisecond, "all buffered records should be delivered before Stop returns") +} + // Test UDP server implementation var ( receivedUDPData [][]byte