Merge pull request #856 from WJL3333/fix_trace_context_cancel_problem

[ISSUE #774] Avoid cancel context too early cause trace data fail.

Author: duhenglucky

consumer/push_consumer.go

@@ -1160,7 +1160,7 @@ func (pc *pushConsumer) consumeMessageOrderly(pq *processQueue, mq *primitive.Me
 				})
 			}
 
-			// jsut put consumeResult in consumerMessageCtx
+			// just put consumeResult in consumerMessageCtx
 			//interval = time.Now().Sub(beginTime)
 			//consumeReult := SuccessReturn
 			//if interval > pc.option.ConsumeTimeout {

internal/trace.go

@@ -463,8 +463,8 @@ func (td *traceDispatcher) sendTraceDataByMQ(keySet Keyset, regionID string, dat
 
 	var req = td.buildSendRequest(mq, msg)
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
 	err := td.cli.InvokeAsync(ctx, addr, req, func(command *remote.RemotingCommand, e error) {
+		cancel()
 		resp := primitive.NewSendResult()
 		if e != nil {
 			rlog.Info("send trace data error.", map[string]interface{}{
@@ -479,6 +479,7 @@ func (td *traceDispatcher) sendTraceDataByMQ(keySet Keyset, regionID string, dat
 		}
 	})
 	if err != nil {
+		cancel()
 		rlog.Info("send trace data error when invoke", map[string]interface{}{
 			rlog.LogKeyUnderlayError: err,
 		})

producer/producer.go

@@ -359,8 +359,9 @@ func (p *defaultProducer) sendAsync(ctx context.Context, msg *primitive.Message,
 		return errors.Errorf("topic=%s route info not found", mq.Topic)
 	}
 
-	ctx, _ = context.WithTimeout(ctx, 3*time.Second)
-	return p.client.InvokeAsync(ctx, addr, p.buildSendRequest(mq, msg), func(command *remote.RemotingCommand, err error) {
+	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
+	err := p.client.InvokeAsync(ctx, addr, p.buildSendRequest(mq, msg), func(command *remote.RemotingCommand, err error) {
+		cancel()
 		resp := primitive.NewSendResult()
 		if err != nil {
 			h(ctx, nil, err)
@@ -369,6 +370,12 @@ func (p *defaultProducer) sendAsync(ctx context.Context, msg *primitive.Message,
 			h(ctx, resp, nil)
 		}
 	})
+
+	if err != nil {
+		cancel()
+	}
+
+	return err
 }
 
 func (p *defaultProducer) SendOneWay(ctx context.Context, msgs ...*primitive.Message) error {