Skip to content

Commit 4ac7d7e

Browse files
lpcoxlpcox
authored
Refactor duplicated tracing and HTTP server setup paths in cmd/server packages (#4048)
Duplicate-code analysis identified repeated logic in three areas: OTLP flag registration, tracing provider lifecycle handling, and HTTP server construction. This PR factors those paths into shared helpers while preserving existing behavior in `root` and `proxy` command flows. - **Tracing flags deduplication (`internal/cmd`)** - Added `registerTracingFlags(...)` to centralize OTLP flag registration (`otlp-endpoint`, `otlp-service-name`, `otlp-sample-rate`). - Replaced duplicated flag wiring in: - `internal/cmd/flags_tracing.go` - `internal/cmd/proxy.go` - **Tracing provider lifecycle deduplication (`internal/cmd`)** - Added shared helpers: - `initTracingProviderWithFallback(...)` - `shutdownTracingProviderWithTimeout(...)` - Reused in both command paths: - `internal/cmd/root.go` - `internal/cmd/proxy.go` - This removes parallel fallback/shutdown implementations while keeping existing warning behavior and noop fallback semantics. - **HTTP server creation deduplication (`internal/server`)** - Added `newHTTPServer(addr, handler)` in `internal/server/http_server.go`. - Replaced repeated `&http.Server{Addr, Handler}` construction in: - `internal/server/transport.go` - `internal/server/routed.go` - **Focused test additions** - Added `internal/cmd/tracing_helpers_test.go` to verify tracing flag defaults and binding behavior. - Added `internal/server/http_server_test.go` to verify shared HTTP server helper wiring. ```go // shared tracing flag registration registerTracingFlags( cmd.Flags(), &proxyOTLPEndpoint, &proxyOTLPService, &proxyOTLPSampleRate, "OTLP HTTP endpoint for trace export (e.g. http://localhost:4318). Tracing is disabled when empty.", "Service name reported in traces.", "Fraction of traces to sample and export (0.0–1.0).", ) ``` > [!WARNING] > > <details> > <summary>Firewall rules blocked me from connecting to one or more addresses (expand for details)</summary> > > #### I tried to connect to the following addresses, but was blocked by firewall rules: > > - `example.com` > - Triggering command: `/tmp/go-build3127029850/b510/launcher.test /tmp/go-build3127029850/b510/launcher.test -test.testlogfile=/tmp/go-build3127029850/b510/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true rotocol/go-sdk@v1.5.0/auth/auth.go rotocol/go-sdk@v1.5.0/auth/authorization_code.go x_amd64/vet -p github.com/tetra-atomic -lang=go1.24 x_amd64/vet 4815�� g_.a -I x_amd64/vet --gdwarf-5 go-sdk/mcp -o x_amd64/vet` (dns block) > - Triggering command: `/tmp/go-build3648569470/b514/launcher.test /tmp/go-build3648569470/b514/launcher.test -test.testlogfile=/tmp/go-build3648569470/b514/testlog.txt -test.paniconexit0 -test.timeout=10m0s` (dns block) > - `invalid-host-that-does-not-exist-12345.com` > - Triggering command: `/tmp/go-build3127029850/b492/config.test /tmp/go-build3127029850/b492/config.test -test.testlogfile=/tmp/go-build3127029850/b492/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true mt-go@v0.1.8/format.go mt-go@v0.1.8/parse.go x_amd64/vet --gdwarf-5 backoff -o x_amd64/vet 4815�� g_.a pkg/mod/go.opent-ifaceassert x_amd64/vet --gdwarf-5 --64 -o x_amd64/vet` (dns block) > - Triggering command: `/tmp/go-build3648569470/b496/config.test /tmp/go-build3648569470/b496/config.test -test.testlogfile=/tmp/go-build3648569470/b496/testlog.txt -test.paniconexit0 -test.timeout=10m0s` (dns block) > - `nonexistent.local` > - Triggering command: `/tmp/go-build3127029850/b510/launcher.test /tmp/go-build3127029850/b510/launcher.test -test.testlogfile=/tmp/go-build3127029850/b510/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true rotocol/go-sdk@v1.5.0/auth/auth.go rotocol/go-sdk@v1.5.0/auth/authorization_code.go x_amd64/vet -p github.com/tetra-atomic -lang=go1.24 x_amd64/vet 4815�� g_.a -I x_amd64/vet --gdwarf-5 go-sdk/mcp -o x_amd64/vet` (dns block) > - Triggering command: `/tmp/go-build3648569470/b514/launcher.test /tmp/go-build3648569470/b514/launcher.test -test.testlogfile=/tmp/go-build3648569470/b514/testlog.txt -test.paniconexit0 -test.timeout=10m0s` (dns block) > - `slow.example.com` > - Triggering command: `/tmp/go-build3127029850/b510/launcher.test /tmp/go-build3127029850/b510/launcher.test -test.testlogfile=/tmp/go-build3127029850/b510/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true rotocol/go-sdk@v1.5.0/auth/auth.go rotocol/go-sdk@v1.5.0/auth/authorization_code.go x_amd64/vet -p github.com/tetra-atomic -lang=go1.24 x_amd64/vet 4815�� g_.a -I x_amd64/vet --gdwarf-5 go-sdk/mcp -o x_amd64/vet` (dns block) > - Triggering command: `/tmp/go-build3648569470/b514/launcher.test /tmp/go-build3648569470/b514/launcher.test -test.testlogfile=/tmp/go-build3648569470/b514/testlog.txt -test.paniconexit0 -test.timeout=10m0s` (dns block) > - `this-host-does-not-exist-12345.com` > - Triggering command: `/tmp/go-build3127029850/b519/mcp.test /tmp/go-build3127029850/b519/mcp.test -test.testlogfile=/tmp/go-build3127029850/b519/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true .cfg -I x_amd64/vet -I . -imultiarch x_amd64/vet -W .cfg om/segmentio/enc-ifaceassert x_amd64/vet . g/grpc/internal/--version --64 x_amd64/vet` (dns block) > - Triggering command: `/tmp/go-build3648569470/b523/mcp.test /tmp/go-build3648569470/b523/mcp.test -test.testlogfile=/tmp/go-build3648569470/b523/testlog.txt -test.paniconexit0 -test.timeout=10m0s -uns�� -unreachable=false /tmp/go-build3127029850/b025/vet.cfg ache/go/1.25.8/x64/pkg/tool/linux_amd64/compile rnal/filetype/bubash ache/go/1.25.8/x/usr/bin/runc x_amd64/compile ache/go/1.25.8/x64/pkg/tool/linux_amd64/compile -uns�� 7029850/b512/_pkg_.a /tmp/go-build3127029850/b122/vet.cfg 7029850/b512=&gt; .go b/gh-aw-mcpg/int/usr/bin/runc x_amd64/compile /opt/hostedtoolcache/go/1.25.8/xorigin` (dns block) > > If you need me to access, download, or install something from one of these locations, you can either: > > - Configure [Actions setup steps](https://gh.io/copilot/actions-setup-steps) to set up my environment, which run before the firewall is enabled > - Add the appropriate URLs or hosts to the custom allowlist in this repository's [Copilot coding agent settings](https://github.com/github/gh-aw-mcpg/settings/copilot/coding_agent) (admins only) > > </details>
2 parents 54e5559 + 9894603 commit 4ac7d7e

10 files changed

Lines changed: 157 additions & 45 deletions

File tree

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ require (
1212
require (
1313
github.com/itchyny/gojq v0.12.19
1414
github.com/santhosh-tekuri/jsonschema/v5 v5.3.1
15+
github.com/spf13/pflag v1.0.9
1516
github.com/stretchr/testify v1.11.1
1617
github.com/tetratelabs/wazero v1.11.0
1718
go.opentelemetry.io/otel v1.43.0
@@ -34,7 +35,6 @@ require (
3435
github.com/pmezard/go-difflib v1.0.0 // indirect
3536
github.com/segmentio/asm v1.1.3 // indirect
3637
github.com/segmentio/encoding v0.5.4 // indirect
37-
github.com/spf13/pflag v1.0.9 // indirect
3838
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
3939
go.opentelemetry.io/auto/sdk v1.2.1 // indirect
4040
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.43.0 // indirect

‎internal/cmd/flags_tracing.go‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,6 @@ package cmd
33
// Tracing-related flags for OpenTelemetry OTLP trace export.
44

55
import (
6-
"github.com/github/gh-aw-mcpg/internal/config"
7-
"github.com/github/gh-aw-mcpg/internal/envutil"
86
"github.com/spf13/cobra"
97
)
108

@@ -17,11 +15,9 @@ var (
1715

1816
func init() {
1917
RegisterFlag(func(cmd *cobra.Command) {
20-
cmd.Flags().StringVar(&otlpEndpoint, "otlp-endpoint", envutil.GetEnvString("OTEL_EXPORTER_OTLP_ENDPOINT", ""),
21-
"OTLP HTTP endpoint for trace export (e.g. http://localhost:4318). Defaults from OTEL_EXPORTER_OTLP_ENDPOINT when set. Tracing is disabled when empty.")
22-
cmd.Flags().StringVar(&otlpServiceName, "otlp-service-name", envutil.GetEnvString("OTEL_SERVICE_NAME", config.DefaultTracingServiceName),
23-
"Service name reported in traces. Defaults from OTEL_SERVICE_NAME when set.")
24-
cmd.Flags().Float64Var(&otlpSampleRate, "otlp-sample-rate", config.DefaultTracingSampleRate,
18+
registerTracingFlags(cmd.Flags(), &otlpEndpoint, &otlpServiceName, &otlpSampleRate,
19+
"OTLP HTTP endpoint for trace export (e.g. http://localhost:4318). Defaults from OTEL_EXPORTER_OTLP_ENDPOINT when set. Tracing is disabled when empty.",
20+
"Service name reported in traces. Defaults from OTEL_SERVICE_NAME when set.",
2521
"Fraction of traces to sample and export (0.0–1.0). Default 1.0 samples everything.")
2622
})
2723
}

‎internal/cmd/proxy.go‎

Lines changed: 14 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package cmd
22

33
import (
4-
"context"
54
"crypto/tls"
65
"fmt"
76
"log"
@@ -12,14 +11,12 @@ import (
1211
"path/filepath"
1312
"strings"
1413
"syscall"
15-
"time"
1614

1715
"github.com/github/gh-aw-mcpg/internal/config"
1816
"github.com/github/gh-aw-mcpg/internal/difc"
1917
"github.com/github/gh-aw-mcpg/internal/envutil"
2018
"github.com/github/gh-aw-mcpg/internal/logger"
2119
"github.com/github/gh-aw-mcpg/internal/proxy"
22-
"github.com/github/gh-aw-mcpg/internal/tracing"
2320
"github.com/spf13/cobra"
2421
)
2522

@@ -121,11 +118,9 @@ Local usage:
121118
cmd.Flags().StringVar(&proxyTLSDir, "tls-dir", "", "Directory for TLS certificates (default: <log-dir>/proxy-tls)")
122119
cmd.Flags().StringSliceVar(&proxyTrustedBots, "trusted-bots", nil, "Additional trusted bot usernames (comma-separated, extends built-in list)")
123120
cmd.Flags().StringSliceVar(&proxyTrustedUsers, "trusted-users", nil, "User logins that receive approved integrity (comma-separated)")
124-
cmd.Flags().StringVar(&proxyOTLPEndpoint, "otlp-endpoint", envutil.GetEnvString("OTEL_EXPORTER_OTLP_ENDPOINT", ""),
125-
"OTLP HTTP endpoint for trace export (e.g. http://localhost:4318). Tracing is disabled when empty.")
126-
cmd.Flags().StringVar(&proxyOTLPService, "otlp-service-name", envutil.GetEnvString("OTEL_SERVICE_NAME", config.DefaultTracingServiceName),
127-
"Service name reported in traces.")
128-
cmd.Flags().Float64Var(&proxyOTLPSampleRate, "otlp-sample-rate", config.DefaultTracingSampleRate,
121+
registerTracingFlags(cmd.Flags(), &proxyOTLPEndpoint, &proxyOTLPService, &proxyOTLPSampleRate,
122+
"OTLP HTTP endpoint for trace export (e.g. http://localhost:4318). Tracing is disabled when empty.",
123+
"Service name reported in traces.",
129124
"Fraction of traces to sample and export (0.0–1.0).")
130125

131126
// Only require --guard-wasm when no baked-in guard is available
@@ -165,17 +160,18 @@ func runProxy(cmd *cobra.Command, args []string) error {
165160
SampleRate: &proxyOTLPSampleRate,
166161
}
167162
}
168-
tracingProvider, err := tracing.InitProvider(ctx, tracingCfg)
169-
if err != nil {
170-
log.Printf("Warning: failed to initialize tracing provider: %v", err)
171-
tracingProvider, _ = tracing.InitProvider(ctx, nil)
172-
}
163+
tracingProvider := initTracingProviderWithFallback(
164+
ctx,
165+
tracingCfg,
166+
"failed to initialize tracing provider: %v",
167+
func(format string, args ...any) {
168+
log.Printf("Warning: "+format, args...)
169+
},
170+
)
173171
defer func() {
174-
shutdownCtx, cancelTracing := context.WithTimeout(context.Background(), 5*time.Second)
175-
defer cancelTracing()
176-
if err := tracingProvider.Shutdown(shutdownCtx); err != nil {
177-
log.Printf("Warning: tracing provider shutdown error: %v", err)
178-
}
172+
shutdownTracingProviderWithTimeout(tracingProvider, func(format string, args ...any) {
173+
log.Printf("Warning: "+format, args...)
174+
})
179175
}()
180176
if tracingCfg != nil {
181177
log.Printf("OpenTelemetry tracing enabled for proxy: endpoint=%s, service=%s", proxyOTLPEndpoint, proxyOTLPService)

‎internal/cmd/root.go‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -311,18 +311,18 @@ func run(cmd *cobra.Command, args []string) error {
311311
if cfg.Gateway != nil {
312312
tracingCfg = cfg.Gateway.Tracing
313313
}
314-
tracingProvider, err := tracing.InitProvider(ctx, tracingCfg)
315-
if err != nil {
316-
logger.StartupWarn("Failed to initialize tracing provider: %v", err)
317-
// Non-fatal: continue without tracing
318-
tracingProvider, _ = tracing.InitProvider(ctx, nil)
319-
}
314+
tracingProvider := initTracingProviderWithFallback(
315+
ctx,
316+
tracingCfg,
317+
"Failed to initialize tracing provider: %v",
318+
func(format string, args ...any) {
319+
logger.StartupWarn(format, args...)
320+
},
321+
)
320322
defer func() {
321-
shutdownCtxTracing, cancelTracing := context.WithTimeout(context.Background(), 5*time.Second)
322-
defer cancelTracing()
323-
if err := tracingProvider.Shutdown(shutdownCtxTracing); err != nil {
324-
log.Printf("Warning: tracing provider shutdown error: %v", err)
325-
}
323+
shutdownTracingProviderWithTimeout(tracingProvider, func(format string, args ...any) {
324+
log.Printf("Warning: "+format, args...)
325+
})
326326
}()
327327

328328
// Apply W3C parent context from configured traceId/spanId (spec §4.1.3.6).

‎internal/cmd/tracing_helpers.go‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package cmd
2+
3+
import (
4+
"context"
5+
"time"
6+
7+
"github.com/github/gh-aw-mcpg/internal/config"
8+
"github.com/github/gh-aw-mcpg/internal/envutil"
9+
"github.com/github/gh-aw-mcpg/internal/tracing"
10+
"github.com/spf13/pflag"
11+
)
12+
13+
func registerTracingFlags(flags *pflag.FlagSet, endpoint *string, serviceName *string, sampleRate *float64, endpointUsage string, serviceUsage string, sampleUsage string) {
14+
flags.StringVar(endpoint, "otlp-endpoint", envutil.GetEnvString("OTEL_EXPORTER_OTLP_ENDPOINT", ""),
15+
endpointUsage)
16+
flags.StringVar(serviceName, "otlp-service-name", envutil.GetEnvString("OTEL_SERVICE_NAME", config.DefaultTracingServiceName),
17+
serviceUsage)
18+
flags.Float64Var(sampleRate, "otlp-sample-rate", config.DefaultTracingSampleRate,
19+
sampleUsage)
20+
}
21+
22+
func initTracingProviderWithFallback(
23+
ctx context.Context,
24+
tracingCfg *config.TracingConfig,
25+
initWarningFormat string,
26+
warnf func(format string, args ...any),
27+
) *tracing.Provider {
28+
tracingProvider, err := tracing.InitProvider(ctx, tracingCfg)
29+
if err != nil {
30+
warnf(initWarningFormat, err)
31+
tracingProvider, _ = tracing.InitProvider(ctx, nil)
32+
}
33+
34+
return tracingProvider
35+
}
36+
37+
func shutdownTracingProviderWithTimeout(tracingProvider *tracing.Provider, warnf func(format string, args ...any)) {
38+
shutdownCtxTracing, cancelTracing := context.WithTimeout(context.Background(), 5*time.Second)
39+
defer cancelTracing()
40+
41+
if err := tracingProvider.Shutdown(shutdownCtxTracing); err != nil {
42+
warnf("tracing provider shutdown error: %v", err)
43+
}
44+
}

‎internal/cmd/tracing_helpers_test.go‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package cmd
2+
3+
import (
4+
"testing"
5+
6+
"github.com/spf13/cobra"
7+
"github.com/stretchr/testify/assert"
8+
"github.com/stretchr/testify/require"
9+
10+
"github.com/github/gh-aw-mcpg/internal/config"
11+
)
12+
13+
func TestRegisterTracingFlags_DefaultsFromEnv(t *testing.T) {
14+
t.Setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://collector:4318")
15+
t.Setenv("OTEL_SERVICE_NAME", "test-service")
16+
17+
cmd := &cobra.Command{Use: "test"}
18+
19+
var endpoint string
20+
var service string
21+
var sampleRate float64
22+
23+
registerTracingFlags(
24+
cmd.Flags(),
25+
&endpoint,
26+
&service,
27+
&sampleRate,
28+
"endpoint help",
29+
"service help",
30+
"sample help",
31+
)
32+
33+
actualEndpoint, err := cmd.Flags().GetString("otlp-endpoint")
34+
require.NoError(t, err)
35+
assert.Equal(t, "http://collector:4318", actualEndpoint)
36+
37+
actualService, err := cmd.Flags().GetString("otlp-service-name")
38+
require.NoError(t, err)
39+
assert.Equal(t, "test-service", actualService)
40+
41+
actualSampleRate, err := cmd.Flags().GetFloat64("otlp-sample-rate")
42+
require.NoError(t, err)
43+
assert.Equal(t, config.DefaultTracingSampleRate, actualSampleRate)
44+
45+
err = cmd.ParseFlags([]string{
46+
"--otlp-endpoint=http://override:4318",
47+
"--otlp-service-name=override-service",
48+
"--otlp-sample-rate=0.25",
49+
})
50+
require.NoError(t, err)
51+
assert.Equal(t, "http://override:4318", endpoint)
52+
assert.Equal(t, "override-service", service)
53+
assert.Equal(t, 0.25, sampleRate)
54+
}

‎internal/server/http_server.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
package server
2+
3+
import "net/http"
4+
5+
func newHTTPServer(addr string, handler http.Handler) *http.Server {
6+
return &http.Server{
7+
Addr: addr,
8+
Handler: handler,
9+
}
10+
}

‎internal/server/http_server_test.go‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package server
2+
3+
import (
4+
"net/http"
5+
"testing"
6+
7+
"github.com/stretchr/testify/assert"
8+
"github.com/stretchr/testify/require"
9+
)
10+
11+
func TestNewHTTPServer(t *testing.T) {
12+
handler := http.NewServeMux()
13+
14+
server := newHTTPServer("127.0.0.1:1234", handler)
15+
require.NotNil(t, server)
16+
assert.Equal(t, "127.0.0.1:1234", server.Addr)
17+
assert.Same(t, handler, server.Handler)
18+
}

‎internal/server/routed.go‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -174,10 +174,7 @@ func CreateHTTPServerForRoutedMode(addr string, unifiedServer *UnifiedServer, ap
174174
log.Printf("Registered route: %s", route)
175175
}
176176

177-
return &http.Server{
178-
Addr: addr,
179-
Handler: mux,
180-
}
177+
return newHTTPServer(addr, mux)
181178
}
182179

183180
// createFilteredServer creates an MCP server that only exposes tools for a specific backend

‎internal/server/transport.go‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,5 @@ func CreateHTTPServerForMCP(addr string, unifiedServer *UnifiedServer, apiKey st
4848
mux.Handle("/mcp/", finalHandler)
4949
mux.Handle("/mcp", finalHandler)
5050

51-
return &http.Server{
52-
Addr: addr,
53-
Handler: mux,
54-
}
51+
return newHTTPServer(addr, mux)
5552
}

0 commit comments

Comments
 (0)