Handle latency with response

Signed-off-by: Jie WU <[email protected]>

Author: courageJ

Diff for: pkg/ext-proc/handlers/response.go

@@ -1,6 +1,9 @@
 package handlers
 
 import (
+	"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"
 	klog "k8s.io/klog/v2"
@@ -63,6 +66,18 @@ func (s *Server) HandleResponseHeaders(reqCtx *RequestContext, req *extProcPb.Pr
 }*/
 func (s *Server) HandleResponseBody(reqCtx *RequestContext, req *extProcPb.ProcessingRequest) (*extProcPb.ProcessingResponse, error) {
 	klog.V(3).Info("Processing HandleResponseBody")
+	body := req.Request.(*extProcPb.ProcessingRequest_ResponseBody)
+
+	res := Response{}
+	if err := json.Unmarshal(body.ResponseBody.Body, &res); err != nil {
+		return nil, fmt.Errorf("unmarshaling response body: %v", err)
+	}
+	reqCtx.Response = res
+	// ResponseComplete is to indicate the response is complete. In non-streaming
+	// case, it will be set to be true once the response is processed; in
+	//streaming case, it will be set to be true once the last chunk is processed.
+	reqCtx.ResponseComplete = true
+	klog.V(3).Infof("Response: %+v", res)
 
 	resp := &extProcPb.ProcessingResponse{
 		Response: &extProcPb.ProcessingResponse_ResponseBody{

Diff for: pkg/ext-proc/handlers/server.go

@@ -93,7 +93,7 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 		case *extProcPb.ProcessingRequest_ResponseBody:
 			resp, err = s.HandleResponseBody(reqCtx, req)
 			reqCtx.ResponseCompleteTimestamp = time.Now()
-			if err == nil {
+			if err == nil && reqCtx.ResponseComplete {
 				metrics.RecordRequestLatencies(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.RequestReceivedTimestamp, reqCtx.ResponseCompleteTimestamp)
 			}
 			klog.V(3).Infof("Request context after HandleResponseBody: %+v", reqCtx)
@@ -138,4 +138,5 @@ type RequestContext struct {
 	ResponseCompleteTimestamp time.Time
 	RequestSize int
 	Response  Response
+	ResponseComplete bool
 }

Diff for: pkg/manifests/ext_proc.yaml

@@ -108,7 +108,6 @@ spec:
         request:
           body: Buffered
         response:
-          body: Streamed
       # The timeouts are likely not needed here. We can experiment with removing/tuning them slowly.
       # The connection limits are more important and will cause the opaque: ext_proc_gRPC_error_14 error in Envoy GW if not configured correctly. 
       messageTimeout: 1000s