fix: Add Wrapper and Keepalive for CTraceContext ids (#32746)

See also #32742

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

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()

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()))

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
 }