fix(coderd/database): improve pubsub closure and context cancellation (#7993)

coderd/database/pubsub_test.go

@@ -45,11 +45,11 @@ func TestPubsub(t *testing.T) {
 		event := "test"
 		data := "testing"
 		messageChannel := make(chan []byte)
-		cancelFunc, err = pubsub.Subscribe(event, func(ctx context.Context, message []byte) {
+		unsub, err := pubsub.Subscribe(event, func(ctx context.Context, message []byte) {
 			messageChannel <- message
 		})
 		require.NoError(t, err)
-		defer cancelFunc()
+		defer unsub()
 		go func() {
 			err = pubsub.Publish(event, []byte(data))
 			assert.NoError(t, err)
@@ -72,6 +72,91 @@ func TestPubsub(t *testing.T) {
 		defer pubsub.Close()
 		cancelFunc()
 	})
+
+	t.Run("NotClosedOnCancelContext", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		connectionURL, closePg, err := postgres.Open()
+		require.NoError(t, err)
+		defer closePg()
+		db, err := sql.Open("postgres", connectionURL)
+		require.NoError(t, err)
+		defer db.Close()
+		pubsub, err := database.NewPubsub(ctx, db, connectionURL)
+		require.NoError(t, err)
+		defer pubsub.Close()
+
+		// Provided context must only be active during NewPubsub, not after.
+		cancel()
+
+		event := "test"
+		data := "testing"
+		messageChannel := make(chan []byte)
+		unsub, err := pubsub.Subscribe(event, func(_ context.Context, message []byte) {
+			messageChannel <- message
+		})
+		require.NoError(t, err)
+		defer unsub()
+		go func() {
+			err = pubsub.Publish(event, []byte(data))
+			assert.NoError(t, err)
+		}()
+		message := <-messageChannel
+		assert.Equal(t, string(message), data)
+	})
+
+	t.Run("ClosePropagatesContextCancellationToSubscription", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+		connectionURL, closePg, err := postgres.Open()
+		require.NoError(t, err)
+		defer closePg()
+		db, err := sql.Open("postgres", connectionURL)
+		require.NoError(t, err)
+		defer db.Close()
+		pubsub, err := database.NewPubsub(ctx, db, connectionURL)
+		require.NoError(t, err)
+		defer pubsub.Close()
+
+		event := "test"
+		done := make(chan struct{})
+		called := make(chan struct{})
+		unsub, err := pubsub.Subscribe(event, func(subCtx context.Context, _ []byte) {
+			defer close(done)
+			select {
+			case <-subCtx.Done():
+				assert.Fail(t, "context should not be canceled")
+			default:
+			}
+			close(called)
+			select {
+			case <-subCtx.Done():
+			case <-ctx.Done():
+				assert.Fail(t, "timeout waiting for sub context to be canceled")
+			}
+		})
+		require.NoError(t, err)
+		defer unsub()
+
+		go func() {
+			err := pubsub.Publish(event, nil)
+			assert.NoError(t, err)
+		}()
+
+		select {
+		case <-called:
+		case <-ctx.Done():
+			require.Fail(t, "timeout waiting for handler to be called")
+		}
+		err = pubsub.Close()
+		require.NoError(t, err)
+
+		select {
+		case <-done:
+		case <-ctx.Done():
+			require.Fail(t, "timeout waiting for handler to finish")
+		}
+	})
 }
 
 func TestPubsub_ordering(t *testing.T) {