Use the request's context to allocate ts (#22633)

Signed-off-by: yah01 <yang.cen@zilliz.com>

internal/proxy/channels_time_ticker.go

@@ -90,7 +90,7 @@ func (ticker *channelsTimeTickerImpl) initCurrents(current Timestamp) {
 }
 
 func (ticker *channelsTimeTickerImpl) tick() error {
-	now, err := ticker.tso.AllocOne()
+	now, err := ticker.tso.AllocOne(ticker.ctx)
 	if err != nil {
 		log.Warn("Proxy channelsTimeTickerImpl failed to get ts from tso", zap.Error(err))
 		return err
@@ -168,7 +168,7 @@ func (ticker *channelsTimeTickerImpl) tickLoop() {
 func (ticker *channelsTimeTickerImpl) start() error {
 	ticker.initStatistics()
 
-	current, err := ticker.tso.AllocOne()
+	current, err := ticker.tso.AllocOne(ticker.ctx)
 	if err != nil {
 		return err
 	}

internal/proxy/interface_def.go

@@ -25,7 +25,7 @@ import (
 // use interface tsoAllocator to keep other components testable
 // include: channelsTimeTickerImpl, baseTaskQueue, taskScheduler
 type tsoAllocator interface {
-	AllocOne() (Timestamp, error)
+	AllocOne(ctx context.Context) (Timestamp, error)
 }
 
 // use timestampAllocatorInterface to keep other components testable

internal/proxy/mock_test.go

@@ -67,7 +67,7 @@ func newMockTimestampAllocatorInterface() timestampAllocatorInterface {
 type mockTsoAllocator struct {
 }
 
-func (tso *mockTsoAllocator) AllocOne() (Timestamp, error) {
+func (tso *mockTsoAllocator) AllocOne(ctx context.Context) (Timestamp, error) {
 	return Timestamp(time.Now().UnixNano()), nil
 }
 

internal/proxy/proxy.go

@@ -207,7 +207,7 @@ func (node *Proxy) Init() error {
 	log.Debug("create id allocator done", zap.String("role", typeutil.ProxyRole), zap.Int64("ProxyID", paramtable.GetNodeID()))
 
 	log.Debug("create timestamp allocator", zap.String("role", typeutil.ProxyRole), zap.Int64("ProxyID", paramtable.GetNodeID()))
-	tsoAllocator, err := newTimestampAllocator(node.ctx, node.rootCoord, paramtable.GetNodeID())
+	tsoAllocator, err := newTimestampAllocator(node.rootCoord, paramtable.GetNodeID())
 	if err != nil {
 		log.Warn("failed to create timestamp allocator",
 			zap.Error(err),

internal/proxy/task_scheduler.go

@@ -169,7 +169,7 @@ func (queue *baseTaskQueue) Enqueue(t task) error {
 		return err
 	}
 
-	ts, err := queue.tsoAllocatorIns.AllocOne()
+	ts, err := queue.tsoAllocatorIns.AllocOne(t.TraceCtx())
 	if err != nil {
 		return err
 	}

internal/proxy/timestamp.go

@@ -32,24 +32,23 @@ import (
 
 // timestampAllocator implements tsoAllocator.
 type timestampAllocator struct {
-	ctx    context.Context
 	tso    timestampAllocatorInterface
 	peerID UniqueID
 }
 
 // newTimestampAllocator creates a new timestampAllocator
-func newTimestampAllocator(ctx context.Context, tso timestampAllocatorInterface, peerID UniqueID) (*timestampAllocator, error) {
+func newTimestampAllocator(tso timestampAllocatorInterface, peerID UniqueID) (*timestampAllocator, error) {
 	a := &timestampAllocator{
-		ctx:    ctx,
 		peerID: peerID,
 		tso:    tso,
 	}
 	return a, nil
 }
 
-func (ta *timestampAllocator) alloc(count uint32) ([]Timestamp, error) {
+func (ta *timestampAllocator) alloc(ctx context.Context, count uint32) ([]Timestamp, error) {
 	tr := timerecord.NewTimeRecorder("applyTimestamp")
-	ctx, cancel := context.WithTimeout(ta.ctx, 5*time.Second)
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
 	req := &rootcoordpb.AllocTimestampRequest{
 		Base: commonpbutil.NewMsgBase(
 			commonpbutil.WithMsgType(commonpb.MsgType_RequestTSO),
@@ -61,7 +60,6 @@ func (ta *timestampAllocator) alloc(count uint32) ([]Timestamp, error) {
 
 	resp, err := ta.tso.AllocTimestamp(ctx, req)
 	defer func() {
-		cancel()
 		metrics.ProxyApplyTimestampLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Observe(float64(tr.ElapseSpan().Milliseconds()))
 	}()
 
@@ -81,8 +79,8 @@ func (ta *timestampAllocator) alloc(count uint32) ([]Timestamp, error) {
 }
 
 // AllocOne allocates a timestamp.
-func (ta *timestampAllocator) AllocOne() (Timestamp, error) {
+func (ta *timestampAllocator) AllocOne(ctx context.Context) (Timestamp, error) {
-	ret, err := ta.alloc(1)
+	ret, err := ta.alloc(ctx, 1)
 	if err != nil {
 		return 0, err
 	}

internal/proxy/timestamp_test.go

@@ -27,11 +27,10 @@ import (
 )
 
 func TestNewTimestampAllocator(t *testing.T) {
-	ctx := context.Background()
 	tso := newMockTimestampAllocatorInterface()
 	peerID := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
 
-	tsAllocator, err := newTimestampAllocator(ctx, tso, peerID)
+	tsAllocator, err := newTimestampAllocator(tso, peerID)
 	assert.Nil(t, err)
 	assert.NotNil(t, tsAllocator)
 }
@@ -41,12 +40,12 @@ func TestTimestampAllocator_alloc(t *testing.T) {
 	tso := newMockTimestampAllocatorInterface()
 	peerID := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
 
-	tsAllocator, err := newTimestampAllocator(ctx, tso, peerID)
+	tsAllocator, err := newTimestampAllocator(tso, peerID)
 	assert.Nil(t, err)
 	assert.NotNil(t, tsAllocator)
 
 	count := rand.Uint32()%100 + 1
-	ret, err := tsAllocator.alloc(count)
+	ret, err := tsAllocator.alloc(ctx, count)
 	assert.Nil(t, err)
 	assert.Equal(t, int(count), len(ret))
 }
@@ -56,10 +55,10 @@ func TestTimestampAllocator_AllocOne(t *testing.T) {
 	tso := newMockTimestampAllocatorInterface()
 	peerID := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
 
-	tsAllocator, err := newTimestampAllocator(ctx, tso, peerID)
+	tsAllocator, err := newTimestampAllocator(tso, peerID)
 	assert.Nil(t, err)
 	assert.NotNil(t, tsAllocator)
 
-	_, err = tsAllocator.AllocOne()
+	_, err = tsAllocator.AllocOne(ctx)
 	assert.Nil(t, err)
 }