Register funcs to support environment

Author: thugrock7

instrumentation/opentelemetry/init.go

@@ -43,7 +43,7 @@ import (
 
 var (
 	batchTimeout          = time.Duration(200) * time.Millisecond
-	traceProviders        map[string]*sdktrace.TracerProvider
+	traceProviders        map[string]map[string]*sdktrace.TracerProvider
 	globalSampler         sdktrace.Sampler
 	initialized           = false
 	enabled               = false
@@ -306,7 +306,7 @@ func InitWithSpanProcessorWrapper(cfg *config.AgentConfig, wrapper SpanProcessor
 
 	otel.SetTextMapPropagator(makePropagator(cfg.PropagationFormats))
 
-	traceProviders = make(map[string]*sdktrace.TracerProvider)
+	traceProviders = make(map[string]map[string]*sdktrace.TracerProvider)
 	globalSampler = sampler
 	initialized = true
 
@@ -324,14 +324,17 @@ func InitWithSpanProcessorWrapper(cfg *config.AgentConfig, wrapper SpanProcessor
 	return func() {
 		mu.Lock()
 		defer mu.Unlock()
-		for serviceName, tracerProvider := range traceProviders {
-			err := tracerProvider.Shutdown(context.Background())
-			if err != nil {
-				log.Printf("error while shutting down tracer provider: %v\n", err)
+		for environment, serviceTraceProviders := range traceProviders {
+			for _, tracerProvider := range serviceTraceProviders {
+				err := tracerProvider.Shutdown(context.Background())
+				if err != nil {
+					log.Printf("error while shutting down tracer provider: %v\n", err)
+				}
 			}
-			delete(traceProviders, serviceName)
+			delete(traceProviders, environment)
 		}
-		traceProviders = map[string]*sdktrace.TracerProvider{}
+
+		traceProviders = map[string]map[string]*sdktrace.TracerProvider{}
 		err := tp.Shutdown(context.Background())
 		if err != nil {
 			log.Printf("error while shutting down default tracer provider: %v\n", err)
@@ -359,14 +362,14 @@ func createResources(serviceName string, resources map[string]string,
 	return retValues
 }
 
-// RegisterService creates tracerprovider for a new service and returns a func which can be used to create spans and the TracerProvider
-func RegisterService(serviceName string, resourceAttributes map[string]string) (sdk.StartSpan, trace.TracerProvider, error) {
-	return RegisterServiceWithSpanProcessorWrapper(serviceName, resourceAttributes, nil, versionInfoAttributes)
+// RegisterService creates tracerprovider for a new service in given environment and returns a func which can be used to create spans and the TracerProvider
+func RegisterService(environment string, serviceName string, resourceAttributes map[string]string) (sdk.StartSpan, trace.TracerProvider, error) {
+	return RegisterServiceWithSpanProcessorWrapper(environment, serviceName, resourceAttributes, nil, versionInfoAttributes)
 }
 
-// RegisterServiceWithSpanProcessorWrapper creates a tracerprovider for a new service with a wrapper over opentelemetry span processor
+// RegisterServiceWithSpanProcessorWrapper creates a tracerprovider for a new service in given environment with a wrapper over opentelemetry span processor
 // and returns a func which can be used to create spans and the TracerProvider
-func RegisterServiceWithSpanProcessorWrapper(serviceName string, resourceAttributes map[string]string,
+func RegisterServiceWithSpanProcessorWrapper(environment string, serviceName string, resourceAttributes map[string]string,
 	wrapper SpanProcessorWrapper, versionInfoAttrs []attribute.KeyValue) (sdk.StartSpan, trace.TracerProvider, error) {
 	mu.Lock()
 	defer mu.Unlock()
@@ -378,8 +381,10 @@ func RegisterServiceWithSpanProcessorWrapper(serviceName string, resourceAttribu
 		return NoopStartSpan, noop.NewTracerProvider(), nil
 	}
 
-	if _, ok := traceProviders[serviceName]; ok {
-		return nil, noop.NewTracerProvider(), fmt.Errorf("service %v already initialized", serviceName)
+	if stp, ok := traceProviders[environment]; ok {
+		if _, present := stp[serviceName]; present {
+			return nil, noop.NewTracerProvider(), fmt.Errorf("environment %v service %v already initialized", environment, serviceName)
+		}
 	}
 
 	exporter, err := exporterFactory()
@@ -408,7 +413,11 @@ func RegisterServiceWithSpanProcessorWrapper(serviceName string, resourceAttribu
 		sdktrace.WithResource(resources),
 	)
 
-	traceProviders[serviceName] = tp
+	if _, ok := traceProviders[environment]; !ok {
+		traceProviders[environment] = make(map[string]*sdktrace.TracerProvider)
+	}
+
+	traceProviders[environment][serviceName] = tp
 	return startSpan(func() trace.TracerProvider {
 		return tp
 	}), tp, nil

instrumentation/opentelemetry/init_test.go

@@ -38,19 +38,38 @@ func ExampleRegisterService() {
 	shutdown := Init(cfg)
 	defer shutdown()
 
-	_, _, err := RegisterService("custom_service", map[string]string{"test1": "val1"})
+	_, _, err := RegisterService("custom_env", "custom_service", map[string]string{"test1": "val1"})
 	if err != nil {
 		log.Fatalf("Error while initializing service: %v", err)
 	}
 }
 
+func TestMultiRegisterService(t *testing.T) {
+	cfg := config.Load()
+	cfg.ServiceName = config.String("my_example_svc")
+	cfg.DataCapture.HttpHeaders.Request = config.Bool(true)
+	cfg.Reporting.TraceReporterType = config.TraceReporterType_LOGGING
+
+	shutdown := Init(cfg)
+	defer shutdown()
+
+	_, _, err := RegisterService("custom_env", "custom_service", map[string]string{"test1": "val1"})
+	assert.Nil(t, err)
+
+	_, _, err = RegisterService("custom_env", "custom_service_2", map[string]string{"test2": "val2"})
+	assert.Nil(t, err)
+
+	_, _, err = RegisterService("custom_env_1", "custom_service", map[string]string{"test3": "val3"})
+	assert.Nil(t, err)
+}
+
 func TestInitDisabledAgent(t *testing.T) {
 	cfg := config.Load()
 	cfg.Enabled = config.Bool(false)
 	shutdown := Init(cfg)
 	defer shutdown()
 
-	startSpan, tp, err := RegisterService("test_service", nil)
+	startSpan, tp, err := RegisterService("custom_env", "test_service", nil)
 	require.NoError(t, err)
 	assert.Equal(t, noop.NewTracerProvider(), tp)
 	_, s, _ := startSpan(context.Background(), "test_span", nil)
@@ -80,7 +99,7 @@ func TestOtlpService(t *testing.T) {
 	shutdown := Init(cfg)
 	defer shutdown()
 
-	startSpan, tp, err := RegisterService("custom_service", map[string]string{"test1": "val1"})
+	startSpan, tp, err := RegisterService("custom_env", "custom_service", map[string]string{"test1": "val1"})
 	_, s, _ := startSpan(context.Background(), "test_span", nil)
 	assert.False(t, s.IsNoop())
 	assert.NotEqual(t, noop.NewTracerProvider(), tp)
@@ -102,7 +121,7 @@ func TestGrpcLoadBalancingConfig(t *testing.T) {
 	defer shutdown()
 
 	assert.Equal(t, resolver.GetDefaultScheme(), "dns")
-	_, tp, err := RegisterService("custom_service", map[string]string{"test1": "val1"})
+	_, tp, err := RegisterService("custom_env", "custom_service", map[string]string{"test1": "val1"})
 	assert.NotEqual(t, noop.NewTracerProvider(), tp)
 	if err != nil {
 		log.Fatalf("Error while initializing service: %v", err)
@@ -165,7 +184,7 @@ func TestMultipleTraceProviders(t *testing.T) {
 	_, _, spanEnder := StartSpan(context.Background(), "example_span", nil)
 	spanEnder()
 
-	startServiceSpan, tp, err := RegisterService("custom_service", map[string]string{"test1": "val1"})
+	startServiceSpan, tp, err := RegisterService("custom_env", "custom_service", map[string]string{"test1": "val1"})
 	assert.NoError(t, err)
 	assert.NotNil(t, startServiceSpan)
 	assert.True(t, initialized)
@@ -208,7 +227,7 @@ func TestMultipleTraceProvidersCallAfterShutdown(t *testing.T) {
 	assert.True(t, initialized)
 	assert.Equal(t, 0, len(traceProviders))
 
-	startServiceSpan, _, err := RegisterService("custom_service", map[string]string{"test1": "val1"})
+	startServiceSpan, _, err := RegisterService("custom_env", "custom_service", map[string]string{"test1": "val1"})
 	assert.NoError(t, err)
 	assert.NotNil(t, startServiceSpan)
 	assert.True(t, initialized)
@@ -379,7 +398,7 @@ func TestInitWithSpanProcessorWrapper(t *testing.T) {
 	assert.Equal(t, 2, wrapper.onEndCount)
 
 	// test wrapper is called for spans created by service trace provider
-	startSpan, _, err := RegisterServiceWithSpanProcessorWrapper("custom_service", map[string]string{"test1": "val1"}, wrapper,
+	startSpan, _, err := RegisterServiceWithSpanProcessorWrapper("custom_env", "custom_service", map[string]string{"test1": "val1"}, wrapper,
 		versionInfoAttributes)
 	if err != nil {
 		log.Fatalf("Error while initializing service: %v", err)