Attach context to client SDK go routines that were missing it (#688)

Author: merlimat

common/batch/batcher_factory.go

@@ -15,8 +15,12 @@
 package batch
 
 import (
+	"context"
+	"fmt"
 	"runtime"
 	"time"
+
+	"github.com/streamnative/oxia/common"
 )
 
 var batcherChannelBufferSize = runtime.GOMAXPROCS(-1)
@@ -26,7 +30,7 @@ type BatcherFactory struct {
 	MaxRequestsPerBatch int
 }
 
-func (b *BatcherFactory) NewBatcher(batchFactory func() Batch) Batcher {
+func (b *BatcherFactory) NewBatcher(ctx context.Context, shard int64, batcherType string, batchFactory func() Batch) Batcher {
 	batcher := &batcherImpl{
 		batchFactory:        batchFactory,
 		callC:               make(chan any, batcherChannelBufferSize),
@@ -35,7 +39,10 @@ func (b *BatcherFactory) NewBatcher(batchFactory func() Batch) Batcher {
 		maxRequestsPerBatch: b.MaxRequestsPerBatch,
 	}
 
-	go batcher.Run()
+	go common.DoWithLabels(ctx, map[string]string{
+		"oxia":  fmt.Sprintf("batcher-%s", batcherType),
+		"shard": fmt.Sprintf("%d", shard),
+	}, batcher.Run)
 
 	return batcher
 }

common/batch/batcher_test.go

@@ -15,6 +15,7 @@
 package batch
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -78,7 +79,7 @@ func TestBatcher(t *testing.T) {
 				Linger:              item.linger,
 				MaxRequestsPerBatch: item.maxSize,
 			}
-			batcher := factory.NewBatcher(batchFactory)
+			batcher := factory.NewBatcher(context.Background(), 1, "test-write", batchFactory)
 			batcher.Add(1)
 
 			if item.closeImmediately {

oxia/async_client_impl.go

@@ -85,10 +85,10 @@ func NewAsyncClient(serviceAddress string, opts ...ClientOption) (AsyncClient, e
 		options:      options,
 		clientPool:   clientPool,
 		shardManager: shardManager,
-		writeBatchManager: batch.NewManager(func(shard *int64) commonbatch.Batcher {
-			return batcherFactory.NewWriteBatcher(shard, options.maxBatchSize)
+		writeBatchManager: batch.NewManager(ctx, func(ctx context.Context, shard *int64) commonbatch.Batcher {
+			return batcherFactory.NewWriteBatcher(ctx, shard, options.maxBatchSize)
 		}),
-		readBatchManager: batch.NewManager(batcherFactory.NewReadBatcher),
+		readBatchManager: batch.NewManager(ctx, batcherFactory.NewReadBatcher),
 		executor:         executor,
 	}
 

oxia/internal/batch/batcher_factory.go

@@ -15,6 +15,7 @@
 package batch
 
 import (
+	"context"
 	"time"
 
 	"github.com/streamnative/oxia/common/batch"
@@ -49,25 +50,25 @@ func NewBatcherFactory(
 	}
 }
 
-func (b *BatcherFactory) NewWriteBatcher(shardId *int64, maxWriteBatchSize int) batch.Batcher {
-	return b.newBatcher(shardId, writeBatchFactory{
+func (b *BatcherFactory) NewWriteBatcher(ctx context.Context, shardId *int64, maxWriteBatchSize int) batch.Batcher {
+	return b.newBatcher(ctx, shardId, "write", writeBatchFactory{
 		execute:        b.Executor.ExecuteWrite,
 		metrics:        b.Metrics,
 		requestTimeout: b.RequestTimeout,
 		maxByteSize:    maxWriteBatchSize,
 	}.newBatch)
 }
 
-func (b *BatcherFactory) NewReadBatcher(shardId *int64) batch.Batcher {
-	return b.newBatcher(shardId, readBatchFactory{
+func (b *BatcherFactory) NewReadBatcher(ctx context.Context, shardId *int64) batch.Batcher {
+	return b.newBatcher(ctx, shardId, "read", readBatchFactory{
 		execute:        b.Executor.ExecuteRead,
 		metrics:        b.Metrics,
 		requestTimeout: b.RequestTimeout,
 	}.newBatch)
 }
 
-func (b *BatcherFactory) newBatcher(shardId *int64, batchFactory func(shardId *int64) batch.Batch) batch.Batcher {
-	return b.NewBatcher(func() batch.Batch {
+func (b *BatcherFactory) newBatcher(ctx context.Context, shardId *int64, batcherType string, batchFactory func(shardId *int64) batch.Batch) batch.Batcher {
+	return b.NewBatcher(ctx, *shardId, batcherType, func() batch.Batch {
 		return batchFactory(shardId)
 	})
 }

oxia/internal/batch/manager.go

@@ -15,23 +15,26 @@
 package batch
 
 import (
+	"context"
 	"sync"
 
 	"go.uber.org/multierr"
 
 	"github.com/streamnative/oxia/common/batch"
 )
 
-func NewManager(batcherFactory func(*int64) batch.Batcher) *Manager {
+func NewManager(ctx context.Context, batcherFactory func(context.Context, *int64) batch.Batcher) *Manager {
 	return &Manager{
+		ctx:            ctx,
 		batcherFactory: batcherFactory,
 		batchers:       make(map[int64]batch.Batcher),
 	}
 }
 
 type Manager struct {
 	sync.RWMutex
-	batcherFactory func(*int64) batch.Batcher
+	ctx            context.Context
+	batcherFactory func(context.Context, *int64) batch.Batcher
 	batchers       map[int64]batch.Batcher
 }
 
@@ -49,7 +52,7 @@ func (m *Manager) Get(shardId int64) batch.Batcher {
 	defer m.Unlock()
 
 	if batcher, ok = m.batchers[shardId]; !ok {
-		batcher = m.batcherFactory(&shardId)
+		batcher = m.batcherFactory(m.ctx, &shardId)
 		m.batchers[shardId] = batcher
 	}
 	return batcher

oxia/internal/batch/manager_test.go

@@ -15,6 +15,7 @@
 package batch
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -42,12 +43,12 @@ func TestManager(t *testing.T) {
 	testBatcher := &testBatcher{}
 
 	newBatcherInvocations := 0
-	batcherFactory := func(*int64) batch.Batcher {
+	batcherFactory := func(context.Context, *int64) batch.Batcher {
 		newBatcherInvocations++
 		return testBatcher
 	}
 
-	manager := NewManager(batcherFactory)
+	manager := NewManager(context.Background(), batcherFactory)
 
 	batcher := manager.Get(shardId)
 	assert.Equal(t, testBatcher, batcher)

oxia/internal/executor.go

@@ -133,7 +133,7 @@ func (e *executorImpl) writeStream(shardId *int64) (*streamWrapper, error) {
 		return nil, err
 	}
 
-	sw = newStreamWrapper(stream)
+	sw = newStreamWrapper(*shardId, stream)
 
 	e.Lock()
 	defer e.Unlock()

oxia/internal/write_stream.go

@@ -16,6 +16,7 @@ package internal
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"log/slog"
 	"sync"
@@ -33,14 +34,20 @@ type streamWrapper struct {
 	failed          atomic.Bool
 }
 
-func newStreamWrapper(stream proto.OxiaClient_WriteStreamClient) *streamWrapper {
+func newStreamWrapper(shard int64, stream proto.OxiaClient_WriteStreamClient) *streamWrapper {
 	sw := &streamWrapper{
 		stream:          stream,
 		pendingRequests: nil,
 	}
 
-	go sw.handleResponses()
-	go sw.handleStreamClosed()
+	go common.DoWithLabels(stream.Context(), map[string]string{
+		"oxia":  "write-stream-handle-response",
+		"shard": fmt.Sprintf("%d", shard),
+	}, sw.handleResponses)
+	go common.DoWithLabels(stream.Context(), map[string]string{
+		"oxia":  "write-stream-handle-stream-closed",
+		"shard": fmt.Sprintf("%d", shard),
+	}, sw.handleStreamClosed)
 	return sw
 }