diff --git a/client/internal/netflow/conntrack/conntrack.go b/client/internal/netflow/conntrack/conntrack.go index 50efefa9c..53619ab8c 100644 --- a/client/internal/netflow/conntrack/conntrack.go +++ b/client/internal/netflow/conntrack/conntrack.go @@ -79,6 +79,7 @@ func (c *ConnTrack) Start(enableCounters bool) error { conn, err := c.dial() if err != nil { + c.RestoreAccounting() return fmt.Errorf("dial conntrack: %w", err) } c.conn = conn @@ -94,6 +95,7 @@ func (c *ConnTrack) Start(enableCounters bool) error { log.Errorf("Error closing conntrack connection: %v", err) } c.conn = nil + c.RestoreAccounting() return fmt.Errorf("start conntrack listener: %w", err) } @@ -252,15 +254,16 @@ func (c *ConnTrack) Close() error { c.started = false + var closeErr error if c.conn != nil { - err := c.conn.Close() + closeErr = c.conn.Close() c.conn = nil + } - c.RestoreAccounting() + c.RestoreAccounting() - if err != nil { - return fmt.Errorf("close conntrack: %w", err) - } + if closeErr != nil { + return fmt.Errorf("close conntrack: %w", closeErr) } return nil diff --git a/client/internal/netflow/conntrack/conntrack_test.go b/client/internal/netflow/conntrack/conntrack_test.go index 696ad6039..1d8ec268a 100644 --- a/client/internal/netflow/conntrack/conntrack_test.go +++ b/client/internal/netflow/conntrack/conntrack_test.go @@ -106,8 +106,12 @@ func TestStopDuringReconnectBackoff(t *testing.T) { // Trigger an error so the receiver enters reconnect. mock.errChan <- assert.AnError - // Give the goroutine time to enter the reconnect backoff wait. - time.Sleep(500 * time.Millisecond) + // Wait for the error handler to close the old listener before calling Stop. + select { + case <-mock.closedCh: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for reconnect to start") + } // Stop while reconnecting. ct.Stop()