adding command name detection and remaining tests

Signed-off-by: TJ Zhang <[email protected]>

Author: tjzhang-BQ

ffi/src/lib.rs

@@ -940,6 +940,7 @@ fn valkey_value_to_command_response(value: Value) -> RedisResult<CommandResponse
 /// * `route_bytes` is an optional array of bytes that will be parsed into a Protobuf `Routes` object. The array must be allocated by the caller and subsequently freed by the caller after this function returns.
 /// * `route_bytes_len` is the number of bytes in `route_bytes`. It must also not be greater than the max value of a signed pointer-sized integer.
 /// * `route_bytes_len` must be 0 if `route_bytes` is null.
+/// * `span_ptr` is a pointer to a span created by [`create_otel_span`]. The span must be valid until the command is finished.
 /// * This function should only be called should with a `client_adapter_ptr` created by [`create_client`], before [`close_client`] was called with the pointer.
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn command(
@@ -974,7 +975,7 @@ pub unsafe extern "C" fn command(
         cmd.arg(command_arg);
     }
     if span_ptr != 0 {
-        cmd.set_span(get_unsafe_span_from_ptr(Some(span_ptr)));
+        cmd.set_span(unsafe { get_unsafe_span_from_ptr(Some(span_ptr)) });
     }
 
     let route = if !route_bytes.is_null() {
@@ -1376,24 +1377,25 @@ pub unsafe extern "C" fn invoke_script(
 /// Creates an OpenTelemetry span with the given name and returns a pointer to the span as u64.
 ///
 /// #Safety
-/// TODO
+/// * `request_type` must be a valid request type.
 #[unsafe(no_mangle)]
-pub unsafe extern "C" fn create_otel_span(name: *const c_char) -> u64 {
-    if name.is_null() {
-        return 0;
-    }
-    let c_str = unsafe { CStr::from_ptr(name) };
-    let name_str = match c_str.to_str() {
-        Ok(s) => s,
-        Err(_) => return 0,
-    };
-    let span = GlideOpenTelemetry::new_span(name_str);
+pub unsafe extern "C" fn create_otel_span(request_type: RequestType) -> u64 {
+    let cmd = request_type
+        .get_command()
+        .expect("Couldn't fetch command type");
+    let cmd_bytes = cmd.command().unwrap();
+    let command_name = std::str::from_utf8(cmd_bytes.as_slice()).unwrap();
+
+    let span = GlideOpenTelemetry::new_span(command_name);
     let arc = Arc::new(span);
     let ptr = Arc::into_raw(arc);
     ptr as u64
 }
 
 /// Drops an OpenTelemetry span given its pointer as u64.
+///
+/// # Safety
+/// * `span_ptr` must be a valid pointer to a span created by [`create_otel_span`].
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn drop_otel_span(span_ptr: u64) {
     if span_ptr == 0 {
@@ -1570,7 +1572,7 @@ fn ensure_tokio_runtime() -> &'static Runtime {
 ///
 /// * `Some(GlideSpan)` - A cloned GlideSpan if the pointer is valid
 /// * `None` - If the pointer is None
-fn get_unsafe_span_from_ptr(command_span: Option<u64>) -> Option<GlideSpan> {
+unsafe fn get_unsafe_span_from_ptr(command_span: Option<u64>) -> Option<GlideSpan> {
     command_span.map(|command_span| unsafe {
         Arc::increment_strong_count(command_span as *const GlideSpan);
         (*Arc::from_raw(command_span as *const GlideSpan)).clone()

go/base_client.go

@@ -287,8 +287,6 @@ func (client *baseClient) executeCommandWithRoute(
 	args []string,
 	route config.Route,
 ) (*C.struct_CommandResponse, error) {
-	//fmt.Println("Starting executeCommandWithRoute===============================")
-	//fmt.Println("otelInstance.ShouldSample()===============================", otelInstance.ShouldSample())
 	// Check if context is already done
 	select {
 	case <-ctx.Done():
@@ -299,14 +297,9 @@ func (client *baseClient) executeCommandWithRoute(
 	// Create span if OpenTelemetry is enabled and sampling is configured
 	var spanPtr uint64
 	if otelInstance != nil && otelInstance.ShouldSample() {
-		// Use the first argument as the command name, or "Batch" if it's a batch command
-		spanName := "Batch"
-		if len(args) > 0 {
-			spanName = args[0]
-		}
-		//fmt.Println("ready to create span===============================")
-		spanPtr = otelInstance.CreateSpan(spanName)
-		//fmt.Println("span created===============================", spanPtr)
+		// Pass the request type to determine the descriptive name of the command
+		// to use as the span name
+		spanPtr = otelInstance.CreateSpan(requestType)
 	}
 	var cArgsPtr *C.uintptr_t = nil
 	var argLengthsPtr *C.ulong = nil
@@ -321,15 +314,13 @@ func (client *baseClient) executeCommandWithRoute(
 		routeProto, err := routeToProtobuf(route)
 		if err != nil {
 			if spanPtr != 0 {
-				//fmt.Println("Dropping span at point 0===============================")
 				otelInstance.DropSpan(spanPtr)
 			}
 			return nil, &errors.RequestError{Msg: "ExecuteCommand failed due to invalid route"}
 		}
 		msg, err := proto.Marshal(routeProto)
 		if err != nil {
 			if spanPtr != 0 {
-				//fmt.Println("Dropping span at point 1===============================")
 				otelInstance.DropSpan(spanPtr)
 			}
 			return nil, err
@@ -350,7 +341,6 @@ func (client *baseClient) executeCommandWithRoute(
 	if client.coreClient == nil {
 		client.mu.Unlock()
 		if spanPtr != 0 {
-			//fmt.Println("Dropping span at point 2===============================")
 			otelInstance.DropSpan(spanPtr)
 		}
 		return nil, &errors.ClosingError{Msg: "ExecuteCommand failed. The client is closed."}
@@ -373,7 +363,6 @@ func (client *baseClient) executeCommandWithRoute(
 	select {
 	case <-ctx.Done():
 		if spanPtr != 0 {
-			//fmt.Println("Dropping span at point 3===============================")
 			otelInstance.DropSpan(spanPtr)
 		}
 		client.mu.Lock()
@@ -391,7 +380,6 @@ func (client *baseClient) executeCommandWithRoute(
 		delete(client.pending, resultChannelPtr)
 	}
 	if spanPtr != 0 {
-		//fmt.Println("Dropping span at point 4===============================")
 		otelInstance.DropSpan(spanPtr)
 	}
 	client.mu.Unlock()

go/integTest/opentelemetry_test.go

@@ -1,3 +1,5 @@
+// Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
+
 package integTest
 
 import (
@@ -18,22 +20,21 @@ import (
 )
 
 const (
-	//otelSpanFile               = "/tmp/spans.json"
 	otelSpanFlushIntervalMs = 100
 	otelSpanTimeoutMs       = 50000
 	validEndpointTraces     = "/tmp/spans.json"
 	validFileEndpointTraces = "file://" + validEndpointTraces
 	validEndpointMetrics    = "https://valid-endpoint/v1/metrics"
 )
 
-type OpenTelemetryIntegrationSuite struct {
+type OpenTelemetryTestSuite struct {
 	GlideTestSuite
 }
 
 // OpenTelemetry standalone tests
 
 // SetupSuite runs once before all tests
-func (suite *OpenTelemetryIntegrationSuite) SetupSuite() {
+func (suite *OpenTelemetryTestSuite) SetupSuite() {
 	// One-time setup for all tests
 	suite.GlideTestSuite.SetupSuite()
 	WrongOpenTelemetryConfig(suite)
@@ -53,12 +54,12 @@ func (suite *OpenTelemetryIntegrationSuite) SetupSuite() {
 }
 
 // TearDownSuite runs once after all tests
-func (suite *OpenTelemetryIntegrationSuite) TearDownSuite() {
+func (suite *OpenTelemetryTestSuite) TearDownSuite() {
 	// One-time cleanup for all tests
 	suite.GlideTestSuite.TearDownSuite()
 }
 
-func (suite *OpenTelemetryIntegrationSuite) TearDownTest() {
+func (suite *OpenTelemetryTestSuite) TearDownTest() {
 	time.Sleep(100 * time.Millisecond)
 	// Remove the span file if it exists
 	if _, err := os.Stat(validEndpointTraces); err == nil {
@@ -70,7 +71,7 @@ func (suite *OpenTelemetryIntegrationSuite) TearDownTest() {
 	suite.GlideTestSuite.TearDownTest()
 }
 
-func WrongOpenTelemetryConfig(suite *OpenTelemetryIntegrationSuite) {
+func WrongOpenTelemetryConfig(suite *OpenTelemetryTestSuite) {
 	// Test wrong traces endpoint
 	cfg := api.OpenTelemetryConfig{
 		Traces: &api.OpenTelemetryTracesConfig{
@@ -185,7 +186,7 @@ func readAndParseSpanFile(path string) (SpanFileData, error) {
 	}, nil
 }
 
-func (suite *OpenTelemetryIntegrationSuite) TestOpenTelemetry_AutomaticSpanLifecycle() {
+func (suite *OpenTelemetryTestSuite) TestOpenTelemetry_AutomaticSpanLifecycle() {
 	suite.runWithSpecificClients(ClientTypeFlag(StandaloneFlag), func(client api.BaseClient) {
 		// Force garbage collection
 		runtime.GC()
@@ -213,7 +214,7 @@ func (suite *OpenTelemetryIntegrationSuite) TestOpenTelemetry_AutomaticSpanLifec
 	})
 }
 
-func (suite *OpenTelemetryIntegrationSuite) TestOpenTelemetry_GlobalConfigNotReinitialize() {
+func (suite *OpenTelemetryTestSuite) TestOpenTelemetry_GlobalConfigNotReinitialize() {
 	suite.runWithSpecificClients(ClientTypeFlag(StandaloneFlag), func(client api.BaseClient) {
 		// Try to initialize OpenTelemetry with wrong endpoint
 		wrongConfig := api.OpenTelemetryConfig{
@@ -239,7 +240,7 @@ func (suite *OpenTelemetryIntegrationSuite) TestOpenTelemetry_GlobalConfigNotRei
 	})
 }
 
-func (suite *OpenTelemetryIntegrationSuite) TestOpenTelemetry_ConcurrentCommandsSpanLifecycle() {
+func (suite *OpenTelemetryTestSuite) TestOpenTelemetry_ConcurrentCommandsSpanLifecycle() {
 	suite.runWithSpecificClients(ClientTypeFlag(StandaloneFlag), func(client api.BaseClient) {
 		// Force garbage collection
 		runtime.GC()
@@ -298,6 +299,152 @@ func (suite *OpenTelemetryIntegrationSuite) TestOpenTelemetry_ConcurrentCommands
 	})
 }
 
-func TestOpenTelemetryIntegrationSuite(t *testing.T) {
-	suite.Run(t, new(OpenTelemetryIntegrationSuite))
+// cluster tests
+func (suite *OpenTelemetryTestSuite) TestOpenTelemetry_ClusterClientMemoryLeak() {
+	suite.runWithSpecificClients(ClientTypeFlag(ClusterFlag), func(client api.BaseClient) {
+		// Force garbage collection
+		runtime.GC()
+
+		// Get initial memory stats
+		var startMem runtime.MemStats
+		runtime.ReadMemStats(&startMem)
+
+		// Execute multiple commands sequentially
+		for i := 0; i < 100; i++ {
+			key := fmt.Sprintf("test_key_%d", i)
+			_, err := client.Set(context.Background(), key, fmt.Sprintf("value_%d", i))
+			require.NoError(suite.T(), err)
+			_, err = client.Get(context.Background(), key)
+			require.NoError(suite.T(), err)
+		}
+
+		// Force garbage collection again
+		runtime.GC()
+
+		// Get final memory stats
+		var endMem runtime.MemStats
+		runtime.ReadMemStats(&endMem)
+
+		// Allow small fluctuations (10% increase)
+		maxAllowedMemory := float64(startMem.HeapAlloc) * 1.1
+		assert.Less(suite.T(), float64(endMem.HeapAlloc), maxAllowedMemory,
+			"Memory usage should not increase significantly")
+	})
+}
+
+func (suite *OpenTelemetryTestSuite) TestOpenTelemetry_ClusterClientSamplingPercentage() {
+	suite.runWithSpecificClients(ClientTypeFlag(ClusterFlag), func(client api.BaseClient) {
+		// Set sampling percentage to 0
+		err := api.GetInstance().SetSamplePercentage(0)
+		require.NoError(suite.T(), err)
+		assert.Equal(suite.T(), int32(0), api.GetInstance().GetSamplePercentage())
+
+		// Wait for any existing spans to be flushed
+		time.Sleep(500 * time.Millisecond)
+
+		// Remove any existing span file
+		if _, err := os.Stat(validEndpointTraces); err == nil {
+			err = os.Remove(validEndpointTraces)
+			require.NoError(suite.T(), err)
+		}
+
+		// Execute commands with 0% sampling
+		for i := 0; i < 100; i++ {
+			_, err := client.Set(context.Background(), "GlideClusterClient_test_percentage_requests_config", "value")
+			require.NoError(suite.T(), err)
+		}
+
+		// Wait for spans to be flushed
+		time.Sleep(500 * time.Millisecond)
+
+		// Verify no spans were exported
+		_, err = os.Stat(validEndpointTraces)
+		assert.True(suite.T(), os.IsNotExist(err), "Span file should not exist with 0% sampling")
+
+		// Set sampling percentage to 100
+		err = api.GetInstance().SetSamplePercentage(100)
+		require.NoError(suite.T(), err)
+
+		// Execute commands with 100% sampling
+		for i := 0; i < 10; i++ {
+			key := fmt.Sprintf("GlideClusterClient_test_percentage_requests_config_%d", i)
+			_, err := client.Get(context.Background(), key)
+			require.NoError(suite.T(), err)
+		}
+
+		// Wait for spans to be flushed
+		time.Sleep(5 * time.Second)
+
+		// Read and verify spans
+		spans, err := readAndParseSpanFile(validEndpointTraces)
+		require.NoError(suite.T(), err)
+
+		// Count Get spans
+		getSpanCount := 0
+		for _, name := range spans.SpanNames {
+			if name == "GET" {
+				getSpanCount++
+			}
+		}
+
+		// Verify we have exactly 10 Get spans
+		assert.Equal(suite.T(), 10, getSpanCount, "Should have exactly 10 Get spans")
+	})
+}
+
+func (suite *OpenTelemetryTestSuite) TestOpenTelemetry_ClusterClientGlobalConfigNotReinitialize() {
+	suite.runWithSpecificClients(ClientTypeFlag(ClusterFlag), func(client api.BaseClient) {
+		// Try to initialize OpenTelemetry with wrong endpoint
+		wrongConfig := api.OpenTelemetryConfig{
+			Traces: &api.OpenTelemetryTracesConfig{
+				Endpoint:         "wrong.endpoint",
+				SamplePercentage: 1,
+			},
+		}
+
+		// The init should not throw error because it can only be initialized once per process
+		err := api.GetInstance().Init(wrongConfig)
+		assert.NoError(suite.T(), err, "OpenTelemetry should not throw error on reinitialization")
+
+		// Execute a command to verify spans are still being exported
+		_, err = client.Set(context.Background(), "GlideClusterClient_test_otel_global_config", "value")
+		require.NoError(suite.T(), err)
+
+		// Wait for spans to be flushed
+		time.Sleep(500 * time.Millisecond)
+
+		// Read spans to verify they're still being exported to the correct endpoint
+		spans, err := readAndParseSpanFile(validEndpointTraces)
+		require.NoError(suite.T(), err)
+		assert.Contains(suite.T(), spans.SpanNames, "SET", "Should find SET span in exported spans")
+	})
+}
+
+func (suite *OpenTelemetryTestSuite) TestOpenTelemetry_ClusterClientMultipleClients() {
+	suite.runWithSpecificClients(ClientTypeFlag(ClusterFlag), func(client1 api.BaseClient) {
+		// Create a second client with the same configuration
+		client2 := suite.clusterClient(suite.defaultClusterClientConfig())
+		defer client2.Close()
+
+		// Execute commands with both clients
+		_, err := client1.Set(context.Background(), "test_key", "value")
+		require.NoError(suite.T(), err)
+		_, err = client2.Get(context.Background(), "test_key")
+		require.NoError(suite.T(), err)
+
+		// Wait for spans to be flushed
+		time.Sleep(5 * time.Second)
+
+		// Read and verify spans
+		spans, err := readAndParseSpanFile(validEndpointTraces)
+		require.NoError(suite.T(), err)
+
+		// Verify both SET and GET spans exist
+		assert.Contains(suite.T(), spans.SpanNames, "SET", "Should find SET span in exported spans")
+		assert.Contains(suite.T(), spans.SpanNames, "GET", "Should find GET span in exported spans")
+	})
+}
+
+func TestOpenTelemetryTestSuite(t *testing.T) {
+	suite.Run(t, new(OpenTelemetryTestSuite))
 }

go/opentelemetry.go

@@ -151,13 +151,11 @@ func (o *OpenTelemetry) SetSamplePercentage(percentage int32) error {
 }
 
 // CreateSpan creates a new OpenTelemetry span with the given name and returns a pointer to the span.
-func (o *OpenTelemetry) CreateSpan(name string) uint64 {
-	if !o.IsInitialized() || name == "" {
+func (o *OpenTelemetry) CreateSpan(requestType C.RequestType) uint64 {
+	if !o.IsInitialized() {
 		return 0
 	}
-	cName := C.CString(name)
-	defer C.free(unsafe.Pointer(cName))
-	return uint64(C.create_otel_span(cName))
+	return uint64(C.create_otel_span(uint32(requestType)))
 }
 
 // DropSpan drops an OpenTelemetry span given its pointer.