diff --git a/internal/querynodev2/segments/load_index_info.go b/internal/querynodev2/segments/load_index_info.go index f35c63292e..c5c1572475 100644 --- a/internal/querynodev2/segments/load_index_info.go +++ b/internal/querynodev2/segments/load_index_info.go @@ -26,6 +26,7 @@ import "C" import ( "context" + "runtime" "unsafe" "github.com/pingcap/log" @@ -223,7 +224,8 @@ func (li *LoadIndexInfo) appendIndexData(ctx context.Context, indexKeys []string status = C.AppendIndexV3(li.cLoadIndexInfo) } else { traceCtx := ParseCTraceContext(ctx) - status = C.AppendIndexV2(traceCtx, li.cLoadIndexInfo) + status = C.AppendIndexV2(traceCtx.ctx, li.cLoadIndexInfo) + runtime.KeepAlive(traceCtx) } return nil, nil }).Await() diff --git a/internal/querynodev2/segments/segment.go b/internal/querynodev2/segments/segment.go index db25aa2785..4e405d35da 100644 --- a/internal/querynodev2/segments/segment.go +++ b/internal/querynodev2/segments/segment.go @@ -29,6 +29,7 @@ import ( "context" "fmt" "io" + "runtime" "strconv" "strings" "unsafe" @@ -557,13 +558,14 @@ func (s *LocalSegment) Search(ctx context.Context, searchReq *SearchRequest) (*S var status C.CStatus GetSQPool().Submit(func() (any, error) { tr := timerecord.NewTimeRecorder("cgoSearch") - status = C.Search(traceCtx, + status = C.Search(traceCtx.ctx, s.ptr, searchReq.plan.cSearchPlan, searchReq.cPlaceholderGroup, C.uint64_t(searchReq.mvccTimestamp), &searchResult.cSearchResult, ) + runtime.KeepAlive(traceCtx) metrics.QueryNodeSQSegmentLatencyInCore.WithLabelValues(fmt.Sprint(paramtable.GetNodeID()), metrics.SearchLabel).Observe(float64(tr.ElapseSpan().Milliseconds())) return nil, nil }).Await() @@ -601,13 +603,14 @@ func (s *LocalSegment) Retrieve(ctx context.Context, plan *RetrievePlan) (*segco GetSQPool().Submit(func() (any, error) { ts := C.uint64_t(plan.Timestamp) tr := timerecord.NewTimeRecorder("cgoRetrieve") - status = C.Retrieve(traceCtx, + status = C.Retrieve(traceCtx.ctx, s.ptr, plan.cRetrievePlan, ts, &retrieveResult.cRetrieveResult, C.int64_t(maxLimitSize), C.bool(plan.ignoreNonPk)) + runtime.KeepAlive(traceCtx) metrics.QueryNodeSQSegmentLatencyInCore.WithLabelValues(fmt.Sprint(paramtable.GetNodeID()), metrics.QueryLabel).Observe(float64(tr.ElapseSpan().Milliseconds())) @@ -674,12 +677,13 @@ func (s *LocalSegment) RetrieveByOffsets(ctx context.Context, plan *RetrievePlan var status C.CStatus tr := timerecord.NewTimeRecorder("cgoRetrieveByOffsets") - status = C.RetrieveByOffsets(traceCtx, + status = C.RetrieveByOffsets(traceCtx.ctx, s.ptr, plan.cRetrievePlan, &retrieveResult.cRetrieveResult, (*C.int64_t)(unsafe.Pointer(&offsets[0])), C.int64_t(len(offsets))) + runtime.KeepAlive(traceCtx) metrics.QueryNodeSQSegmentLatencyInCore.WithLabelValues(fmt.Sprint(paramtable.GetNodeID()), metrics.QueryLabel).Observe(float64(tr.ElapseSpan().Milliseconds())) diff --git a/internal/querynodev2/segments/trace.go b/internal/querynodev2/segments/trace.go index 663e67953a..7fb9c565bf 100644 --- a/internal/querynodev2/segments/trace.go +++ b/internal/querynodev2/segments/trace.go @@ -30,15 +30,27 @@ import ( "go.opentelemetry.io/otel/trace" ) +// CTraceContext is the wrapper for `C.CTraceContext` +// it stores the internal C.CTraceContext and +type CTraceContext struct { + traceID trace.TraceID + spanID trace.SpanID + ctx C.CTraceContext +} + // ParseCTraceContext parses tracing span and convert it into `C.CTraceContext`. -func ParseCTraceContext(ctx context.Context) C.CTraceContext { +func ParseCTraceContext(ctx context.Context) *CTraceContext { span := trace.SpanFromContext(ctx) - traceID := span.SpanContext().TraceID() - spanID := span.SpanContext().SpanID() - return C.CTraceContext{ - traceID: (*C.uint8_t)(unsafe.Pointer(&traceID[0])), - spanID: (*C.uint8_t)(unsafe.Pointer(&spanID[0])), + cctx := &CTraceContext{ + traceID: span.SpanContext().TraceID(), + spanID: span.SpanContext().SpanID(), + } + cctx.ctx = C.CTraceContext{ + traceID: (*C.uint8_t)(unsafe.Pointer(&cctx.traceID[0])), + spanID: (*C.uint8_t)(unsafe.Pointer(&cctx.spanID[0])), traceFlags: (C.uint8_t)(span.SpanContext().TraceFlags()), } + + return cctx }