Repository: honeycombio/refinery
Pull request #607 (2023-02-10 20:25:06)

refactor: Change Router to use TraceServer
## Which problem is this PR solving?

- Updates Refinery to use the gRPC server pattern that Shepherd uses.  This will allow us to add support for metrics and logs proxying in the future.

## Short description of the changes

- Adds new TraceServer that implements the `TraceServiceServer` interface.
- Unimplement `TraceServiceServer` for `Router`.



github-actions[bot]: Thank you for contributing to the project! 🎉

We require pull request titles to follow the [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/) and it looks like your proposed title needs to be adjusted.

Make sure to prepend with `feat:`, `fix:`, or another option in the list below.

Once you update the title, this workflow will re-run automatically and validate the updated title.

Details:

```
No release type found in pull request title "Refactor Router to use TraceServer". Add a prefix to indicate what kind of release this pull request corresponds to. For reference, see https://www.conventionalcommits.org/

Available types:
 - maint
 - rel
 - fix
 - feat
 - chore
 - ci
 - docs
 - style
 - refactor
 - perf
 - test
```
<!-- Sticky Pull Request Commentpr-title-lint-error -->

Patch:
diff --git a/route/otlp_trace.go b/route/otlp_trace.go
index 6b75b29..6c71493 100644
--- a/route/otlp_trace.go
+++ b/route/otlp_trace.go
@@ -11,29 +11,39 @@ import (
 	collectortrace "go.opentelemetry.io/proto/otlp/collector/trace/v1"
 )
 
-func (router *Router) postOTLP(w http.ResponseWriter, req *http.Request) {
+func (r *Router) postOTLP(w http.ResponseWriter, req *http.Request) {
 	ri := huskyotlp.GetRequestInfoFromHttpHeaders(req.Header)
 	if err := ri.ValidateTracesHeaders(); err != nil {
 		if errors.Is(err, huskyotlp.ErrInvalidContentType) {
-			router.handlerReturnWithError(w, ErrInvalidContentType, err)
+			r.handlerReturnWithError(w, ErrInvalidContentType, err)
 		} else {
-			router.handlerReturnWithError(w, ErrAuthNeeded, err)
+			r.handlerReturnWithError(w, ErrAuthNeeded, err)
 		}
 		return
 	}
 
 	result, err := huskyotlp.TranslateTraceRequestFromReader(req.Body, ri)
 	if err != nil {
-		router.handlerReturnWithError(w, ErrUpstreamFailed, err)
+		r.handlerReturnWithError(w, ErrUpstreamFailed, err)
 		return
 	}
 
-	if err := processTraceRequest(req.Context(), router, result.Batches, ri.ApiKey); err != nil {
-		router.handlerReturnWithError(w, ErrUpstreamFailed, err)
+	if err := processTraceRequest(req.Context(), r, result.Batches, ri.ApiKey); err != nil {
+		r.handlerReturnWithError(w, ErrUpstreamFailed, err)
 	}
 }
 
-func (router *Router) Export(ctx context.Context, req *collectortrace.ExportTraceServiceRequest) (*collectortrace.ExportTraceServiceResponse, error) {
+type TraceServer struct {
+	router *Router
+	collectortrace.UnimplementedTraceServiceServer
+}
+
+func NewTraceServer(router *Router) *TraceServer {
+	traceServer := TraceServer{router: router}
+	return &traceServer
+}
+
+func (t *TraceServer) Export(ctx context.Context, req *collectortrace.ExportTraceServiceRequest) (*collectortrace.ExportTraceServiceResponse, error) {
 	ri := huskyotlp.GetRequestInfoFromGrpcMetadata(ctx)
 	if err := ri.ValidateTracesHeaders(); err != nil {
 		return nil, huskyotlp.AsGRPCError(err)
@@ -44,7 +54,7 @@ func (router *Router) Export(ctx context.Context, req *collectortrace.ExportTrac
 		return nil, huskyotlp.AsGRPCError(err)
 	}
 
-	if err := processTraceRequest(ctx, router, result.Batches, ri.ApiKey); err != nil {
+	if err := processTraceRequest(ctx, t.router, result.Batches, ri.ApiKey); err != nil {
 		return nil, huskyotlp.AsGRPCError(err)
 	}
 
diff --git a/route/route.go b/route/route.go
index 21c6916..90e1203 100644
--- a/route/route.go
+++ b/route/route.go
@@ -82,9 +82,6 @@ type Router struct {
 	grpcServer *grpc.Server
 	doneWG     sync.WaitGroup
 
-	// used to identify Router as a OTLP TraceServer
-	collectortrace.UnimplementedTraceServiceServer
-
 	environmentCache *environmentCache
 }
 
@@ -229,8 +226,9 @@ func (r *Router) LnS(incomingOrPeer string) {
 				Timeout:               r.Config.GetGRPCTimeout(),
 			}),
 		}
+		traceServer := NewTraceServer(r)
 		r.grpcServer = grpc.NewServer(serverOpts...)
-		collectortrace.RegisterTraceServiceServer(r.grpcServer, r)
+		collectortrace.RegisterTraceServiceServer(r.grpcServer, traceServer)
 		grpc_health_v1.RegisterHealthServer(r.grpcServer, r)
 		go r.grpcServer.Serve(l)
 	}


Test patch:
diff --git a/route/otlp_trace_test.go b/route/otlp_trace_test.go
index caf9bf2..3baac2d 100644
--- a/route/otlp_trace_test.go
+++ b/route/otlp_trace_test.go
@@ -73,7 +73,8 @@ func TestOTLPHandler(t *testing.T) {
 				}},
 			}},
 		}
-		_, err := router.Export(ctx, req)
+		traceServer := NewTraceServer(router)
+		_, err := traceServer.Export(ctx, req)
 		if err != nil {
 			t.Errorf(`Unexpected error: %s`, err)
 		}
@@ -89,7 +90,8 @@ func TestOTLPHandler(t *testing.T) {
 				}},
 			}},
 		}
-		_, err := router.Export(ctx, req)
+		traceServer := NewTraceServer(router)
+		_, err := traceServer.Export(ctx, req)
 		if err != nil {
 			t.Errorf(`Unexpected error: %s`, err)
 		}
@@ -122,7 +124,8 @@ func TestOTLPHandler(t *testing.T) {
 				}},
 			}},
 		}
-		_, err := router.Export(ctx, req)
+		traceServer := NewTraceServer(router)
+		_, err := traceServer.Export(ctx, req)
 		if err != nil {
 			t.Errorf(`Unexpected error: %s`, err)
 		}
@@ -171,7 +174,8 @@ func TestOTLPHandler(t *testing.T) {
 				}},
 			}},
 		}
-		_, err := router.Export(ctx, req)
+		traceServer := NewTraceServer(router)
+		_, err := traceServer.Export(ctx, req)
 		if err != nil {
 			t.Errorf(`Unexpected error: %s`, err)
 		}
@@ -337,7 +341,8 @@ func TestOTLPHandler(t *testing.T) {
 				}},
 			}},
 		}
-		_, err := router.Export(ctx, req)
+		traceServer := NewTraceServer(router)
+		_, err := traceServer.Export(ctx, req)
 		if err != nil {
 			t.Errorf(`Unexpected error: %s`, err)
 		}
@@ -370,7 +375,8 @@ func TestOTLPHandler(t *testing.T) {
 				}},
 			}},
 		}
-		_, err := router.Export(ctx, req)
+		traceServer := NewTraceServer(router)
+		_, err := traceServer.Export(ctx, req)
 		if err != nil {
 			t.Errorf(`Unexpected error: %s`, err)
 		}