[ENH]: add distributed tracing

chromadb/telemetry/opentelemetry/grpc.py

@@ -4,8 +4,6 @@
 import grpc
 from opentelemetry.trace import StatusCode, SpanKind
 
-from chromadb.telemetry.opentelemetry import tracer
-
 
 class _ClientCallDetails(
     collections.namedtuple(
@@ -36,6 +34,8 @@ class OtelInterceptor(
     grpc.StreamStreamClientInterceptor,
 ):
     def _intercept_call(self, continuation, client_call_details, request_or_iterator):
+        from chromadb.telemetry.opentelemetry import tracer
+
         if tracer is None:
             return continuation(client_call_details, request_or_iterator)
         with tracer.start_as_current_span(

go/cmd/logservice/main.go

@@ -8,6 +8,7 @@ import (
 	"github.com/chroma-core/chroma/go/pkg/proto/logservicepb"
 	"github.com/chroma-core/chroma/go/pkg/utils"
 	libs "github.com/chroma-core/chroma/go/shared/libs"
+	"github.com/chroma-core/chroma/go/shared/otel"
 	"github.com/pingcap/log"
 	"github.com/rs/zerolog"
 	"go.uber.org/automaxprocs/maxprocs"
@@ -18,6 +19,7 @@ import (
 
 func main() {
 	ctx := context.Background()
+
 	// Configure logger
 	utils.LogLevel = zerolog.DebugLevel
 	utils.ConfigureLogger()
@@ -26,6 +28,13 @@ func main() {
 	}
 	log.Info("Starting log service")
 	config := configuration.NewLogServiceConfiguration()
+	err := otel.InitTracing(ctx, &otel.TracingConfig{
+		Service:  "log-service",
+		Endpoint: config.OPTL_TRACING_ENDPOINT,
+	})
+	if err != nil {
+		log.Fatal("failed to initialize tracing", zap.Error(err))
+	}
 	conn, err := libs.NewPgConnection(ctx, config)
 	if err != nil {
 		log.Fatal("failed to connect to postgres", zap.Error(err))
@@ -37,7 +46,7 @@ func main() {
 	if err != nil {
 		log.Fatal("failed to listen", zap.Error(err))
 	}
-	s := grpc.NewServer()
+	s := grpc.NewServer(grpc.UnaryInterceptor(otel.ServerGrpcInterceptor))
 	logservicepb.RegisterLogServiceServer(s, server)
 	log.Info("log service started", zap.String("address", listener.Addr().String()))
 	if err := s.Serve(listener); err != nil {

go/pkg/grpcutils/service.go

@@ -1,6 +1,7 @@
 package grpcutils
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"github.com/chroma-core/chroma/go/shared/otel"
@@ -72,6 +73,14 @@ func newDefaultGrpcProvider(name string, grpcConfig *GrpcConfig, registerFunc fu
 		opts = append(opts, grpc.Creds(credentials.NewTLS(tlsConfig)))
 	}
 	opts = append(opts, grpc.UnaryInterceptor(otel.ServerGrpcInterceptor))
+	OPTL_TRACING_ENDPOINT := os.Getenv("OPTL_TRACING_ENDPOINT")
+	if OPTL_TRACING_ENDPOINT == "" {
+		OPTL_TRACING_ENDPOINT = "jaeger:4317"
+	}
+	otel.InitTracing(context.Background(), &otel.TracingConfig{
+		Service:  "sysdb-service",
+		Endpoint: OPTL_TRACING_ENDPOINT,
+	})
 
 	c := &defaultGrpcServer{
 		server: grpc.NewServer(opts...),

go/pkg/log/configuration/config.go

@@ -3,8 +3,9 @@ package configuration
 import "os"
 
 type LogServiceConfiguration struct {
-	PORT         string
-	DATABASE_URL string
+	PORT                  string
+	DATABASE_URL          string
+	OPTL_TRACING_ENDPOINT string
 }
 
 func getEnvWithDefault(key, defaultValue string) string {
@@ -17,7 +18,8 @@ func getEnvWithDefault(key, defaultValue string) string {
 
 func NewLogServiceConfiguration() *LogServiceConfiguration {
 	return &LogServiceConfiguration{
-		PORT:         getEnvWithDefault("PORT", "50051"),
-		DATABASE_URL: getEnvWithDefault("CHROMA_DATABASE_URL", "postgresql://chroma:chroma@postgres.chroma.svc.cluster.local:5432/log"),
+		PORT:                  getEnvWithDefault("PORT", "50051"),
+		DATABASE_URL:          getEnvWithDefault("CHROMA_DATABASE_URL", "postgresql://chroma:chroma@postgres.chroma.svc.cluster.local:5432/log"),
+		OPTL_TRACING_ENDPOINT: getEnvWithDefault("OPTL_TRACING_ENDPOINT", "jaeger:4317"),
 	}
 }

go/shared/otel/main.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/hex"
 	"fmt"
+	"github.com/pingcap/log"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	otelCode "go.opentelemetry.io/otel/codes"
@@ -13,6 +14,7 @@ import (
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
 	"go.opentelemetry.io/otel/trace"
+	"go.uber.org/zap"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
@@ -71,32 +73,35 @@ func ServerGrpcInterceptor(ctx context.Context, req interface{}, info *grpc.Unar
 	var span trace.Span
 	ctx, span = tracer.Start(ctx, "Request "+info.FullMethod)
 	defer span.End()
+	span.SetAttributes(attribute.String("rpc.method", info.FullMethod))
 
 	// Calls the handler
 	h, err := handler(ctx, req)
 	if err != nil {
 		// Handle and log the error.
-		handleError(span, err)
+		handleError(span, info, err)
 		return nil, err
 	}
 
 	// Set the status to OK upon success.
 	span.SetStatus(otelCode.Ok, "ok")
 	span.SetAttributes(attribute.String("rpc.status_code", "ok"))
-	span.SetAttributes(attribute.String("rpc.method", info.FullMethod))
+	log.Info("RPC call", zap.String("method", info.FullMethod), zap.String("status", "ok"))
 
 	return h, nil
 }
 
 // handleError logs and annotates the span with details of the encountered error.
-func handleError(span trace.Span, err error) {
+func handleError(span trace.Span, info *grpc.UnaryServerInfo, err error) {
 	st, _ := status.FromError(err)
 	span.SetStatus(otelCode.Error, "error")
 	span.SetAttributes(
 		attribute.String("rpc.status_code", st.Code().String()),
 		attribute.String("rpc.message", st.Message()),
 		attribute.String("rpc.error", st.Err().Error()),
 	)
+	log.Error("RPC call", zap.String("method", info.FullMethod), zap.String("status", st.Code().String()), zap.String("error", st.Err().Error()), zap.String("message", st.Message()))
+
 }
 
 // decodeMetadataValue safely extracts a value from metadata, allowing for missing keys.

k8s/distributed-chroma/values.yaml

@@ -27,8 +27,11 @@ frontendService:
   memberlistProviderImpl: 'value: "chromadb.segment.impl.distributed.segment_directory.MockMemberlistProvider"'
   logServiceHost: 'value: "logservice.chroma"'
   logServicePort: 'value: "50051"'
-  otherEnvConfig: ''
-
+  otherEnvConfig: |
+      - name: CHROMA_OTEL_COLLECTION_ENDPOINT
+        value: "http://jaeger:4317"
+      - name: CHROMA_OTEL_GRANULARITY
+        value: all
 sysdb:
   image:
     repository: 'local'