[Draft] Support Multiple Backends in Proxy Mode

main.go

@@ -47,7 +47,6 @@ import (
 )
 
 const (
-	serviceName         = "package-registry"
 	version             = "1.32.2"
 	defaultInstanceName = "localhost"
 )
@@ -83,6 +82,7 @@ var (
 
 	featureProxyMode bool
 	proxyTo          string
+	serviceName      = getServiceName()
 
 	defaultConfig = Config{
 		CacheTimeIndex:               10 * time.Second,
@@ -336,6 +336,13 @@ func getHostname() string {
 	return hostname
 }
 
+func getServiceName() string {
+	if name := os.Getenv("ELASTIC_APM_SERVICE_NAME"); name != "" {
+		return name
+	}
+	return "package-registry"
+}
+
 func initMetricsServer(logger *zap.Logger) {
 	if metricsAddress == "" {
 		return
@@ -630,9 +637,47 @@ func getRouter(logger *zap.Logger, options serverOptions) (*mux.Router, error) {
 	if featureProxyMode {
 		logger.Info("Technical preview: Proxy mode is an experimental feature and it may be unstable.")
 	}
+
+	// Default proxy backend and retry timeout
+	proxyTimeout := 2 * time.Second
+	retryMax := 1
+	retryWaitMin := 1 * time.Second
+	retryWaitMax := 3 * time.Second
+
+	// Override the default from an environment variable if any.
+	if v := os.Getenv("EPR_PROXY_TIMEOUT"); v != "" {
+		if d, err := time.ParseDuration(v); err == nil {
+			proxyTimeout = d
+		} else {
+			logger.Warn("invalid EPR_PROXY_TIMEOUT format, using default", zap.String("value", v))
+		}
+	}
+
+	if v := os.Getenv("EPR_PROXY_RETRY_MAX"); v != "" {
+		if i, err := strconv.Atoi(v); err == nil {
+			retryMax = i
+		}
+	}
+
+	if v := os.Getenv("EPR_PROXY_RETRY_WAIT_MIN"); v != "" {
+		if d, err := time.ParseDuration(v); err == nil {
+			retryWaitMin = d
+		}
+	}
+
+	if v := os.Getenv("EPR_PROXY_RETRY_WAIT_MAX"); v != "" {
+		if d, err := time.ParseDuration(v); err == nil {
+			retryWaitMax = d
+		}
+	}
+
 	proxyMode, err := proxymode.NewProxyMode(logger, proxymode.ProxyOptions{
-		Enabled: featureProxyMode,
-		ProxyTo: proxyTo,
+		Enabled:      featureProxyMode,
+		ProxyTo:      strings.Split(proxyTo, ","),
+		Timeout:      proxyTimeout,
+		RetryMax:     retryMax,
+		RetryWaitMin: retryWaitMin,
+		RetryWaitMax: retryWaitMax,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("can't create proxy mode: %w", err)

proxymode/logger.go

@@ -5,46 +5,73 @@
 package proxymode
 
 import (
+	"context"
+	"errors"
+
 	"github.com/hashicorp/go-retryablehttp"
+	"go.elastic.co/apm/module/apmzap/v2"
 	"go.uber.org/zap"
 )
 
 type zapLoggerAdapter struct {
 	target *zap.Logger
+	ctx    context.Context
 }
 
 var _ retryablehttp.LeveledLogger = new(zapLoggerAdapter)
 
-func withZapLoggerAdapter(target *zap.Logger) retryablehttp.LeveledLogger {
+func newZapLoggerAdapter(ctx context.Context, target *zap.Logger) retryablehttp.LeveledLogger {
 	return &zapLoggerAdapter{
 		target: target,
+		ctx:    ctx,
 	}
 }
 
 func (a zapLoggerAdapter) Error(msg string, keysAndValues ...interface{}) {
-	a.target.Error(msg, keysAndValuesAsZapFields(keysAndValues...)...)
+	// Check if the error is an expected context cancellation.
+	isCanceled := false
+	for i := 0; i < len(keysAndValues); i += 2 {
+		if key, ok := keysAndValues[i].(string); ok && key == "error" {
+			if err, ok := keysAndValues[i+1].(error); ok && errors.Is(err, context.Canceled) {
+				isCanceled = true
+				break
+			}
+		}
+	}
+
+	loggerWithContext := a.target.With(apmzap.TraceContext(a.ctx)...)
+
+	// If the error was an expected cancellation, log it as a debug message.
+	if isCanceled {
+		loggerWithContext.Debug(msg, keysAndValuesAsZapFields(keysAndValues...)...)
+		return
+	}
+
+	loggerWithContext.Error(msg, keysAndValuesAsZapFields(keysAndValues...)...)
 }
 
 func (a zapLoggerAdapter) Info(msg string, keysAndValues ...interface{}) {
-	a.target.Info(msg, keysAndValuesAsZapFields(keysAndValues...)...)
+	loggerWithContext := a.target.With(apmzap.TraceContext(a.ctx)...)
+	loggerWithContext.Info(msg, keysAndValuesAsZapFields(keysAndValues...)...)
 }
 
 func (a zapLoggerAdapter) Debug(msg string, keysAndValues ...interface{}) {
-	a.target.Debug(msg, keysAndValuesAsZapFields(keysAndValues...)...)
+	loggerWithContext := a.target.With(apmzap.TraceContext(a.ctx)...)
+	loggerWithContext.Debug(msg, keysAndValuesAsZapFields(keysAndValues...)...)
 }
 
 func (a zapLoggerAdapter) Warn(msg string, keysAndValues ...interface{}) {
-	a.target.Warn(msg, keysAndValuesAsZapFields(keysAndValues...)...)
+	loggerWithContext := a.target.With(apmzap.TraceContext(a.ctx)...)
+	loggerWithContext.Warn(msg, keysAndValuesAsZapFields(keysAndValues...)...)
 }
 
-// keysAndValuesAsZapFields function transforms the LeveledLogger arguments to the zap.Logger interface.
 func keysAndValuesAsZapFields(keysAndValues ...interface{}) []zap.Field {
 	fields := make([]zap.Field, len(keysAndValues)/2)
 	var j int
 	for i := 0; i < len(keysAndValues); i += 2 {
 		key, ok := keysAndValues[i].(string)
 		if !ok {
-			continue // something is wrong with the key, string expected
+			continue
 		}
 		fields[j] = zap.Any(key, keysAndValues[i+1])
 		j++