perf: optimize memory allocations in CoProcess middleware

coprocess/grpc/coprocess_grpc_test.go

@@ -10,6 +10,7 @@
 	"mime/multipart"
 	"net/http"
 	"os"
+	"runtime"
 	"strconv"
 	"strings"
 	"testing"
@@ -638,3 +639,108 @@
 	})
 
 }
+
+func BenchmarkGRPCDispatch_MemoryOverhead(b *testing.B) {
+	ts, cleanupFn := startTestServices(b)
+	b.Cleanup(cleanupFn)
+
+	// Create a session with a large amount of metadata and access rights
+	// to simulate a real-world scenario and exacerbate ProtoSessionState allocations
+	keyID := gateway.CreateSession(ts.Gw, func(s *user.SessionState) {
+		s.MetaData = map[string]interface{}{}
+		for i := 0; i < 100; i++ {
+			s.MetaData[fmt.Sprintf("key_%d", i)] = strings.Repeat("a", 100)
+		}
+		s.AccessRights = map[string]user.AccessDefinition{
+			"1": {
+				APIID:    "1",
+				APIName:  "API 1",
+				Versions: []string{"Default"},
+			},
+		}
+		for i := 0; i < 50; i++ {
+			s.AccessRights[fmt.Sprintf("api_%d", i)] = user.AccessDefinition{
+				APIID:    fmt.Sprintf("api_%d", i),
+				APIName:  fmt.Sprintf("API %d", i),
+				Versions: []string{"Default"},
+			}
+		}
+	})
+	headers := map[string]string{"authorization": keyID}
+	b.Run("Pre Hook with Large Session", func(b *testing.B) {
+		basepath := "/grpc-test-api/"
+		b.ReportAllocs() // This will print B/op and allocs/op
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			ts.Run(b, test.TestCase{
+				Path:    basepath,
+				Method:  http.MethodGet,
+				Code:    http.StatusOK,
+				Headers: headers,
+			})
+		}
+	})
+
+	emptyKeyID := gateway.CreateSession(ts.Gw, func(s *user.SessionState) {})
+	emptyHeaders := map[string]string{"authorization": emptyKeyID}
+
+	b.Run("Pre Hook with Empty Session", func(b *testing.B) {
+		basepath := "/grpc-test-api/"
+		b.ReportAllocs()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			ts.Run(b, test.TestCase{
+				Path:    basepath,
+				Method:  http.MethodGet,
+				Code:    http.StatusOK,
+				Headers: emptyHeaders,
+			})
+		}
+	})
+}
+
+func TestGRPCDispatch_MemoryLeakCheck(t *testing.T) {
+	ts, cleanupFn := startTestServices(t)
+	t.Cleanup(cleanupFn)
+
+	keyID := gateway.CreateSession(ts.Gw, func(s *user.SessionState) {
+		s.MetaData = map[string]interface{}{}
+		for i := 0; i < 100; i++ {
+			s.MetaData[fmt.Sprintf("key_%d", i)] = strings.Repeat("a", 100)
+		}
+		s.AccessRights = map[string]user.AccessDefinition{
+			"1": {
+				APIID:    "1",
+				APIName:  "API 1",
+				Versions: []string{"Default"},
+			},
+		}
+	})
+	headers := map[string]string{"authorization": keyID}
+
+	var m runtime.MemStats
+	runtime.GC() // Clean up before starting
+	runtime.ReadMemStats(&m)
+	initialAlloc := m.Alloc
+
+	// Run a large number of requests
+	for i := 0; i < 5000; i++ {
+		ts.Run(t, test.TestCase{
+			Path:    "/grpc-test-api/",
+			Method:  http.MethodGet,
+			Code:    http.StatusOK,
+			Headers: headers,
+		})
+	}
+
+	runtime.GC() // Force GC to see what is retained
+	runtime.ReadMemStats(&m)
+	finalAlloc := m.Alloc
+
+	t.Logf("Initial Alloc: %d bytes", initialAlloc)
+	t.Logf("Final Alloc: %d bytes", finalAlloc)
+	t.Logf("Difference: %d bytes", int64(finalAlloc)-int64(initialAlloc))
+
+	// If it was a leak, the difference would be huge (e.g., hundreds of MBs).
+	// Since it's just churn, the difference will be relatively small after GC.
+}

gateway/coprocess.go

@@ -87,14 +87,9 @@
 		scheme = "https"
 	}
 	miniRequestObject := &coprocess.MiniRequestObject{
-		Headers:        headers,
-		SetHeaders:     map[string]string{},
-		DeleteHeaders:  []string{},
-		Url:            req.URL.String(),
-		Params:         ProtoMap(req.URL.Query()),
-		AddParams:      map[string]string{},
-		ExtendedParams: ProtoMap(nil),
-		DeleteParams:   []string{},
+		Headers: headers,
+		Url:     req.URL.String(),
+		Params:  ProtoMap(req.URL.Query()),
 		ReturnOverrides: &coprocess.ReturnOverrides{
 			ResponseCode: -1,
 		},
@@ -121,22 +116,20 @@
 		HookType: c.Middleware.HookType,
 	}
 
-	object.Spec = make(map[string]string)
-
 	// Append spec data:
 	if c.Middleware != nil {
 		configDataAsJSON := []byte("{}")
 		if shouldAddConfigData(c.Middleware.Spec) {
 			var err error
 			configDataAsJSON, err = json.Marshal(c.Middleware.Spec.ConfigData)
 			if err != nil {
 				return nil, err
 			}
 		}
 
 		bundleHash, err := c.Middleware.Gw.getHashedBundleName(spec.CustomMiddlewareBundle)
 		if err != nil {
 			return nil, err
 		}
 
 		object.Spec = map[string]string{
@@ -158,11 +151,11 @@
 			object.Metadata = object.Session.Metadata
 		}
 	}
-
 	// Append response data if it's available:
 	if res != nil {
 		resObj := &coprocess.ResponseObject{
-			Headers: make(map[string]string, len(res.Header)),
+			Headers:           make(map[string]string, len(res.Header)),
+			MultivalueHeaders: make([]*coprocess.Header, 0, len(res.Header)),
 		}
 		for k, v := range res.Header {
 			// set univalue header

gateway/coprocess_helpers.go

@@ -91,41 +91,55 @@
 }
 
 // ProtoSessionState takes a standard SessionState and outputs a SessionState object compatible with Protocol Buffers.
 func ProtoSessionState(session *user.SessionState) *coprocess.SessionState {
 
-	accessDefinitions := make(map[string]*coprocess.AccessDefinition, len(session.AccessRights))
-
-	for key, accessDefinition := range session.AccessRights {
-		var allowedUrls []*coprocess.AccessSpec
-		for _, allowedURL := range accessDefinition.AllowedURLs {
-			accessSpec := &coprocess.AccessSpec{
-				Url:     allowedURL.URL,
-				Methods: allowedURL.Methods,
+	var accessDefinitions map[string]*coprocess.AccessDefinition
+	if len(session.AccessRights) > 0 {
+		accessDefinitions = make(map[string]*coprocess.AccessDefinition, len(session.AccessRights))
+		for key, accessDefinition := range session.AccessRights {
+			allowedUrls := make([]*coprocess.AccessSpec, 0, len(accessDefinition.AllowedURLs))
+			for _, allowedURL := range accessDefinition.AllowedURLs {
+				accessSpec := &coprocess.AccessSpec{
+					Url:     allowedURL.URL,
+					Methods: allowedURL.Methods,
+				}
+				allowedUrls = append(allowedUrls, accessSpec)
 			}
-			allowedUrls = append(allowedUrls, accessSpec)
-		}
 
-		accessDefinitions[key] = &coprocess.AccessDefinition{
-			ApiName:     accessDefinition.APIName,
-			ApiId:       accessDefinition.APIID,
-			Versions:    accessDefinition.Versions,
-			AllowedUrls: allowedUrls,
+			accessDefinitions[key] = &coprocess.AccessDefinition{
+				ApiName:     accessDefinition.APIName,
+				ApiId:       accessDefinition.APIID,
+				Versions:    accessDefinition.Versions,
+				AllowedUrls: allowedUrls,
+			}
 		}
 	}
 
-	basicAuthData := &coprocess.BasicAuthData{
-		Password: session.BasicAuthData.Password,
-		Hash:     string(session.BasicAuthData.Hash),
+	var basicAuthData *coprocess.BasicAuthData
+	if session.BasicAuthData.Password != "" || len(session.BasicAuthData.Hash) > 0 {
+		basicAuthData = &coprocess.BasicAuthData{
+			Password: session.BasicAuthData.Password,
+			Hash:     string(session.BasicAuthData.Hash),
+		}
 	}
-	jwtData := &coprocess.JWTData{
-		Secret: session.JWTData.Secret,
+
+	var jwtData *coprocess.JWTData
+	if session.JWTData.Secret != "" {
+		jwtData = &coprocess.JWTData{
+			Secret: session.JWTData.Secret,
+		}
 	}
-	monitor := &coprocess.Monitor{
-		TriggerLimits: session.Monitor.TriggerLimits,
+
+	var monitor *coprocess.Monitor
+	if len(session.Monitor.TriggerLimits) > 0 {
+		monitor = &coprocess.Monitor{
+			TriggerLimits: session.Monitor.TriggerLimits,
+		}
 	}
 
-	metadata := make(map[string]string)
+	var metadata map[string]string
 	if len(session.MetaData) > 0 {
+		metadata = make(map[string]string, len(session.MetaData))
 		for k, v := range session.MetaData {
 			switch v.(type) {
 			case string:
@@ -181,13 +195,14 @@
 
 // ProtoMap is a helper function for maps with string slice values.
 func ProtoMap(inputMap map[string][]string) map[string]string {
-	newMap := make(map[string]string)
-
-	if inputMap != nil {
-		for k, v := range inputMap {
+	if len(inputMap) == 0 {
+		return nil
+	}
+	newMap := make(map[string]string, len(inputMap))
+	for k, v := range inputMap {
+		if len(v) > 0 {
 			newMap[k] = v[0]
 		}
 	}
-
 	return newMap
 }