trickle: Tighten up goroutine cleanup (#3808)

* trickle: Use goleak in unit tests as necessary.

* trickle: Use base context when setting the seq in subscribers.

This way subscriptions will still be properly terminated when the
passed-in context is cancelled.

* trickle: Close pending writers in the publisher as needed.

This prevents them from hanging for a while which triggers goleak.

Author: j0sh

trickle/local_subscriber_test.go

@@ -14,6 +14,7 @@ func TestLocalSubscriber_OverrunSeq(t *testing.T) {
 
 	pub, err := NewTricklePublisher(url)
 	require.Nil(err)
+	defer pub.Close()
 
 	sub := NewLocalSubscriber(server, "testest")
 

trickle/trickle_publisher.go

@@ -126,6 +126,15 @@ func (c *TricklePublisher) preconnect() (*pendingPost, error) {
 }
 
 func (c *TricklePublisher) Close() error {
+
+	// Close any pending writers
+	c.writeLock.Lock()
+	pp := c.pendingPost
+	if pp != nil {
+		pp.writer.Close()
+	}
+	c.writeLock.Unlock()
+
 	req, err := http.NewRequest("DELETE", c.baseURL, nil)
 	if err != nil {
 		return err
@@ -207,6 +216,7 @@ func (p *pendingPost) Write(data io.Reader) (int64, error) {
 	// before writing, check for error from preconnects
 	select {
 	case err := <-errCh:
+		writer.Close()
 		return 0, err
 	default:
 		// no error, continue
@@ -229,6 +239,7 @@ func (p *pendingPost) Write(data io.Reader) (int64, error) {
 	// also prioritize errors over this channel compared to io errors
 	// such as "read/write on closed pipe"
 	if err := <-errCh; err != nil {
+		writer.Close()
 		return n, err
 	}
 

trickle/trickle_subscriber.go

@@ -60,6 +60,7 @@ type TrickleSubscriber struct {
 	url        string
 	mu         sync.Mutex      // Mutex to manage concurrent access
 	pendingGet *http.Response  // Pre-initialized GET request
+	baseCtx    context.Context // base context to use for the next context
 	ctx        context.Context // Parent context to use for pending GETs. This is bad
 	cancelCtx  func()          // cancel the pending GET
 	idx        int             // Segment index to request
@@ -91,6 +92,7 @@ func NewTrickleSubscriber(config TrickleSubscriberConfig) (*TrickleSubscriber, e
 	return &TrickleSubscriber{
 		client:    httpClient(),
 		url:       config.URL,
+		baseCtx:   baseCtx,
 		ctx:       ctx,
 		cancelCtx: cancel,
 		idx:       idx,
@@ -134,7 +136,7 @@ func (c *TrickleSubscriber) SetSeq(seq int) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	c.idx = seq
-	c.ctx, c.cancelCtx = context.WithCancel(context.Background())
+	c.ctx, c.cancelCtx = context.WithCancel(c.baseCtx)
 	c.pendingGet = nil
 	c.preconnectErrorCount = 0
 }

trickle/trickle_test.go

@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 )
 
 func TestTrickle_Close(t *testing.T) {
@@ -23,15 +24,17 @@ func TestTrickle_Close(t *testing.T) {
 	})
 	stop := server.Start()
 	ts := httptest.NewServer(mux)
+	defer goleak.VerifyNone(t)
 	defer ts.Close()
 	defer stop()
 
 	channelURL := ts.URL + "/testest"
 	pub, err := NewTricklePublisher(channelURL)
 	require.Nil(err)
+	defer pub.Close()
 	require.Error(StreamNotFoundErr, pub.Write(bytes.NewReader([]byte("first post"))))
 
-	sub, err := NewTrickleSubscriber(subConfig(channelURL))
+	sub, err := NewTrickleSubscriber(subConfig(t, channelURL))
 	require.Nil(err)
 	sub.SetSeq(0)
 
@@ -50,6 +53,7 @@ func TestTrickle_Close(t *testing.T) {
 	// now recreate pub, should be ok
 	pub, err = NewTricklePublisher(channelURL)
 	require.Nil(err)
+	defer pub.Close()
 
 	// write two segments
 	segs := []string{"first", "second"}
@@ -82,13 +86,15 @@ func TestTrickle_Close(t *testing.T) {
 	require.Error(EOS, pub.Write(bytes.NewReader([]byte("invalid"))))
 
 	// Spinning up a second subscriber should return 404
-	sub2, err := NewTrickleSubscriber(subConfig(channelURL))
+	sub2, err := NewTrickleSubscriber(subConfig(t, channelURL))
 	require.Nil(err)
 	_, err = sub2.Read()
 	require.Error(StreamNotFoundErr, err)
 
 	// Spinning up a second publisher should return 404
 	pub2, err := NewTricklePublisher(channelURL)
+	require.Nil(err)
+	defer pub2.Close()
 	require.Error(StreamNotFoundErr, pub2.Write(bytes.NewReader([]byte("bad post"))))
 }
 
@@ -97,7 +103,8 @@ func TestTrickle_SetSeq(t *testing.T) {
 
 	pub, err := NewTricklePublisher(channelURL)
 	require.Nil(err)
-	sub, err := NewTrickleSubscriber(subConfig(channelURL))
+	defer pub.Close()
+	sub, err := NewTrickleSubscriber(subConfig(t, channelURL))
 	require.Nil(err)
 
 	// give sub preconnect time to latch on
@@ -144,15 +151,17 @@ func TestTrickle_Reset(t *testing.T) {
 	})
 	stop := server.Start()
 	ts := httptest.NewServer(mux)
+	defer goleak.VerifyNone(t)
 	defer ts.Close()
 	defer stop()
 
 	channelURL := ts.URL + "/testest"
 
 	pub, err := NewTricklePublisher(channelURL)
 	require.Nil(err)
+	defer pub.Close()
 
-	sub, err := NewTrickleSubscriber(subConfig(channelURL))
+	sub, err := NewTrickleSubscriber(subConfig(t, channelURL))
 	require.Nil(err)
 	wg := &sync.WaitGroup{}
 
@@ -241,14 +250,15 @@ func TestTrickle_IdleSweep(t *testing.T) {
 	})
 	stop := server.Start()
 	ts := httptest.NewServer(mux)
+	defer goleak.VerifyNone(t)
 	defer ts.Close()
 	defer stop()
 
 	channelURL := ts.URL + "/testest"
 	lp := NewLocalPublisher(server, channelURL, "text/plain")
 	lp.CreateChannel()
 
-	sub, err := NewTrickleSubscriber(subConfig(channelURL))
+	sub, err := NewTrickleSubscriber(subConfig(t, channelURL))
 	require.Nil(err)
 	_, err = sub.Read()
 	require.ErrorIs(err, StreamNotFoundErr)
@@ -282,6 +292,7 @@ func TestTrickle_SetSubStart(t *testing.T) {
 	// 1. Subscribe from the beginning
 	subBeginning, err := NewTrickleSubscriber(TrickleSubscriberConfig{
 		URL: url,
+		Ctx: t.Context(),
 	})
 	require.Nil(err)
 	wg.Add(1)
@@ -302,12 +313,15 @@ func TestTrickle_SetSubStart(t *testing.T) {
 
 	pub, err := NewTricklePublisher(url)
 	require.Nil(err)
+	defer pub.Close()
+
 	require.Nil(pub.Write(bytes.NewReader([]byte("zeroth"))))
 	require.Nil(pub.Write(bytes.NewReader([]byte("first"))))
 
 	// 2. Subscribe from the current seq
 	seq := Current
 	subCurrent, err := NewTrickleSubscriber(TrickleSubscriberConfig{
+		Ctx:   t.Context(),
 		URL:   url,
 		Start: &seq,
 	})
@@ -329,6 +343,7 @@ func TestTrickle_SetSubStart(t *testing.T) {
 	// 3. Subscribe from the next seq
 	seq = Next
 	subNext, err := NewTrickleSubscriber(TrickleSubscriberConfig{
+		Ctx:   t.Context(),
 		URL:   url,
 		Start: &seq,
 	})
@@ -356,6 +371,7 @@ func TestTrickle_SetSubStart(t *testing.T) {
 	// 4. Subscribe from a specific seq
 	seq = 1
 	subSeq, err := NewTrickleSubscriber(TrickleSubscriberConfig{
+		Ctx:   t.Context(),
 		URL:   url,
 		Start: &seq,
 	})
@@ -395,6 +411,7 @@ func makeServerWithServer(t *testing.T) (*require.Assertions, string, *Server) {
 	t.Cleanup(func() {
 		stop()
 		ts.Close()
+		goleak.VerifyNone(t)
 	})
 
 	// create the channel locally on the server
@@ -405,6 +422,6 @@ func makeServerWithServer(t *testing.T) (*require.Assertions, string, *Server) {
 	return require, ts.URL + "/" + chanName, server
 }
 
-func subConfig(url string) TrickleSubscriberConfig {
-	return TrickleSubscriberConfig{URL: url}
+func subConfig(t *testing.T, url string) TrickleSubscriberConfig {
+	return TrickleSubscriberConfig{URL: url, Ctx: t.Context()}
 }