Pass around ctx instead of a logger

tchap · tchap · commit abde09f861ce · 2025-02-15T14:17:26.000+01:00
diff --git a/pkg/ext-proc/handlers/request.go b/pkg/ext-proc/handlers/request.go
@@ -1,15 +1,16 @@
 package handlers
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"strconv"
 
 	configPb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
-	"github.com/go-logr/logr"
 	"google.golang.org/protobuf/types/known/structpb"
+	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/ext-proc/backend"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/ext-proc/scheduling"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/ext-proc/util/logging"
@@ -19,10 +20,11 @@ import (
 // parameter.
 // Envoy sends the request body to ext proc before sending the request to the backend server.
 func (s *Server) HandleRequestBody(
-	logger logr.Logger,
+	ctx context.Context,
 	reqCtx *RequestContext,
 	req *extProcPb.ProcessingRequest,
 ) (*extProcPb.ProcessingResponse, error) {
+	logger := log.FromContext(ctx)
 	loggerVerbose := logger.V(logutil.VERBOSE)
 	loggerVerbose.Info("Handling request body")
 
@@ -76,7 +78,7 @@ func (s *Server) HandleRequestBody(
 		loggerVerbose.Info("Updated request body marshalled", "body", string(requestBody))
 	}
 
-	targetPod, err := s.scheduler.Schedule(logger, llmReq)
+	targetPod, err := s.scheduler.Schedule(ctx, llmReq)
 	if err != nil {
 		return nil, fmt.Errorf("failed to find target pod: %w", err)
 	}
@@ -141,13 +143,13 @@ func (s *Server) HandleRequestBody(
 }
 
 func HandleRequestHeaders(
-	logger logr.Logger,
+	ctx context.Context,
 	reqCtx *RequestContext,
 	req *extProcPb.ProcessingRequest,
 ) *extProcPb.ProcessingResponse {
 	r := req.Request
 	h := r.(*extProcPb.ProcessingRequest_RequestHeaders)
-	logger.Info("Handling request headers", "headers", h)
+	log.FromContext(ctx).Info("Handling request headers", "headers", h)
 
 	resp := &extProcPb.ProcessingResponse{
 		Response: &extProcPb.ProcessingResponse_RequestHeaders{
diff --git a/pkg/ext-proc/handlers/response.go b/pkg/ext-proc/handlers/response.go
@@ -1,22 +1,23 @@
 package handlers
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 
 	configPb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
-	"github.com/go-logr/logr"
+	"sigs.k8s.io/controller-runtime/pkg/log"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/ext-proc/util/logging"
 )
 
 // HandleResponseHeaders processes response headers from the backend model server.
 func (s *Server) HandleResponseHeaders(
-	logger logr.Logger,
+	ctx context.Context,
 	reqCtx *RequestContext,
 	req *extProcPb.ProcessingRequest,
 ) (*extProcPb.ProcessingResponse, error) {
-	loggerVerbose := logger.V(logutil.VERBOSE)
+	loggerVerbose := log.FromContext(ctx).V(logutil.VERBOSE)
 	loggerVerbose.Info("Processing ResponseHeaders")
 	h := req.Request.(*extProcPb.ProcessingRequest_ResponseHeaders)
 	loggerVerbose.Info("Headers before", "headers", h)
@@ -71,10 +72,11 @@ func (s *Server) HandleResponseHeaders(
     }
 }*/
 func (s *Server) HandleResponseBody(
-	logger logr.Logger,
+	ctx context.Context,
 	reqCtx *RequestContext,
 	req *extProcPb.ProcessingRequest,
 ) (*extProcPb.ProcessingResponse, error) {
+	logger := log.FromContext(ctx)
 	loggerVerbose := logger.V(logutil.VERBOSE)
 	loggerVerbose.Info("Processing HandleResponseBody")
 	body := req.Request.(*extProcPb.ProcessingRequest_ResponseBody)
diff --git a/pkg/ext-proc/handlers/response_test.go b/pkg/ext-proc/handlers/response_test.go
@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"context"
 	"testing"
 
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
@@ -35,7 +36,7 @@ const (
 )
 
 func TestHandleResponseBody(t *testing.T) {
-	logger := logutil.NewTestLogger()
+	ctx := logutil.NewTestLoggerIntoContext(context.Background())
 
 	tests := []struct {
 		name    string
@@ -73,7 +74,7 @@ func TestHandleResponseBody(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			server := &Server{}
 			reqCtx := &RequestContext{}
-			_, err := server.HandleResponseBody(logger, reqCtx, &extProcPb.ProcessingRequest{Request: test.req})
+			_, err := server.HandleResponseBody(ctx, reqCtx, &extProcPb.ProcessingRequest{Request: test.req})
 			if err != nil {
 				if !test.wantErr {
 					t.Fatalf("HandleResponseBody returned unexpected error: %v, want %v", err, test.wantErr)
diff --git a/pkg/ext-proc/handlers/server.go b/pkg/ext-proc/handlers/server.go
@@ -1,12 +1,12 @@
 package handlers
 
 import (
+	"context"
 	"io"
 	"time"
 
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	envoyTypePb "github.com/envoyproxy/go-control-plane/envoy/type/v3"
-	"github.com/go-logr/logr"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"sigs.k8s.io/controller-runtime/pkg/log"
@@ -38,7 +38,7 @@ type Server struct {
 }
 
 type Scheduler interface {
-	Schedule(logger logr.Logger, b *scheduling.LLMRequest) (targetPod backend.Pod, err error)
+	Schedule(ctx context.Context, b *scheduling.LLMRequest) (targetPod backend.Pod, err error)
 }
 
 // PodProvider is an interface to provide set of pods in the backend and information such as metrics.
@@ -83,23 +83,23 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 		switch v := req.Request.(type) {
 		case *extProcPb.ProcessingRequest_RequestHeaders:
 			reqCtx.RequestReceivedTimestamp = time.Now()
-			resp = HandleRequestHeaders(logger, reqCtx, req)
+			resp = HandleRequestHeaders(ctx, reqCtx, req)
 			loggerVerbose.Info("Request context after HandleRequestHeaders", "context", reqCtx)
 		case *extProcPb.ProcessingRequest_RequestBody:
-			resp, err = s.HandleRequestBody(logger, reqCtx, req)
+			resp, err = s.HandleRequestBody(ctx, reqCtx, req)
 			if err == nil {
 				metrics.RecordRequestCounter(reqCtx.Model, reqCtx.ResolvedTargetModel)
 				metrics.RecordRequestSizes(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.RequestSize)
 			}
 			loggerVerbose.Info("Request context after HandleRequestBody", "context", reqCtx)
 		case *extProcPb.ProcessingRequest_ResponseHeaders:
-			resp, err = s.HandleResponseHeaders(logger, reqCtx, req)
+			resp, err = s.HandleResponseHeaders(ctx, reqCtx, req)
 			loggerVerbose.Info("Request context after HandleResponseHeaders", "context", reqCtx)
 		case *extProcPb.ProcessingRequest_ResponseBody:
-			resp, err = s.HandleResponseBody(logger, reqCtx, req)
+			resp, err = s.HandleResponseBody(ctx, reqCtx, req)
 			if err == nil && reqCtx.ResponseComplete {
 				reqCtx.ResponseCompleteTimestamp = time.Now()
-				metrics.RecordRequestLatencies(logger, reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.RequestReceivedTimestamp, reqCtx.ResponseCompleteTimestamp)
+				metrics.RecordRequestLatencies(ctx, reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.RequestReceivedTimestamp, reqCtx.ResponseCompleteTimestamp)
 				metrics.RecordResponseSizes(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.ResponseSize)
 				metrics.RecordInputTokens(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.Response.Usage.PromptTokens)
 				metrics.RecordOutputTokens(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.Response.Usage.CompletionTokens)
diff --git a/pkg/ext-proc/metrics/metrics.go b/pkg/ext-proc/metrics/metrics.go
@@ -1,12 +1,13 @@
 package metrics
 
 import (
+	"context"
 	"sync"
 	"time"
 
-	"github.com/go-logr/logr"
 	compbasemetrics "k8s.io/component-base/metrics"
 	"k8s.io/component-base/metrics/legacyregistry"
+	"sigs.k8s.io/controller-runtime/pkg/log"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/ext-proc/util/logging"
 )
 
@@ -144,9 +145,9 @@ func RecordRequestSizes(modelName, targetModelName string, reqSize int) {
 }
 
 // RecordRequestLatencies records duration of request.
-func RecordRequestLatencies(logger logr.Logger, modelName, targetModelName string, received time.Time, complete time.Time) bool {
+func RecordRequestLatencies(ctx context.Context, modelName, targetModelName string, received time.Time, complete time.Time) bool {
 	if !complete.After(received) {
-		logger.V(logutil.DEFAULT).Error(nil, "Request latency values are invalid",
+		log.FromContext(ctx).V(logutil.DEFAULT).Error(nil, "Request latency values are invalid",
 			"modelName", modelName, "targetModelName", targetModelName, "completeTime", complete, "receivedTime", received)
 		return false
 	}
diff --git a/pkg/ext-proc/metrics/metrics_test.go b/pkg/ext-proc/metrics/metrics_test.go
@@ -1,6 +1,7 @@
 package metrics
 
 import (
+	"context"
 	"os"
 	"testing"
 	"time"
@@ -91,7 +92,7 @@ func TestRecordRequestCounterandSizes(t *testing.T) {
 }
 
 func TestRecordRequestLatencies(t *testing.T) {
-	logger := logutil.NewTestLogger()
+	ctx := logutil.NewTestLoggerIntoContext(context.Background())
 	timeBaseline := time.Now()
 	type requests struct {
 		modelName       string
@@ -150,7 +151,7 @@ func TestRecordRequestLatencies(t *testing.T) {
 	for _, scenario := range scenarios {
 		t.Run(scenario.name, func(t *testing.T) {
 			for _, req := range scenario.reqs {
-				success := RecordRequestLatencies(logger, req.modelName, req.targetModelName, req.receivedTime, req.completeTime)
+				success := RecordRequestLatencies(ctx, req.modelName, req.targetModelName, req.receivedTime, req.completeTime)
 				if success == scenario.invalid {
 					t.Errorf("got record success(%v), but the request expects invalid(%v)", success, scenario.invalid)
 				}
diff --git a/pkg/ext-proc/scheduling/scheduler.go b/pkg/ext-proc/scheduling/scheduler.go
@@ -2,12 +2,14 @@
 package scheduling
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 
 	"github.com/go-logr/logr"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/ext-proc/backend"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/ext-proc/util/logging"
 )
@@ -110,8 +112,8 @@ type PodMetricsProvider interface {
 }
 
 // Schedule finds the target pod based on metrics and the requested lora adapter.
-func (s *Scheduler) Schedule(logger logr.Logger, req *LLMRequest) (targetPod backend.Pod, err error) {
-	logger = logger.WithValues("request", req)
+func (s *Scheduler) Schedule(ctx context.Context, req *LLMRequest) (targetPod backend.Pod, err error) {
+	logger := log.FromContext(ctx).WithValues("request", req)
 	logger.V(logutil.VERBOSE).Info("Scheduling a request", "metrics", s.podMetricsProvider.AllPodMetrics())
 	pods, err := s.filter.Filter(logger, req, s.podMetricsProvider.AllPodMetrics())
 	if err != nil || len(pods) == 0 {
diff --git a/pkg/ext-proc/util/logging/logger.go b/pkg/ext-proc/util/logging/logger.go
@@ -1,12 +1,20 @@
 package logging
 
 import (
+	"context"
+
 	"github.com/go-logr/logr"
 	uberzap "go.uber.org/zap"
+	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 )
 
 // NewTestLogger creates a new Zap logger using the dev mode.
 func NewTestLogger() logr.Logger {
 	return zap.New(zap.UseDevMode(true), zap.RawZapOpts(uberzap.AddCaller()))
 }
+
+// NewTestLoggerIntoContext creates a new Zap logger using the dev mode and inserts it into the given context.
+func NewTestLoggerIntoContext(ctx context.Context) context.Context {
+	return log.IntoContext(ctx, zap.New(zap.UseDevMode(true), zap.RawZapOpts(uberzap.AddCaller())))
+}
