fix: CGO Memory leak issue in GO Feature server (#4291)

EXPEbdodla · web-flow · commit 43e198f6945c · 2024-06-28T20:58:35.000-04:00
diff --git a/go/internal/feast/server/http_server.go b/go/internal/feast/server/http_server.go
@@ -4,14 +4,16 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"net/http"
+	"time"
+
 	"github.com/feast-dev/feast/go/internal/feast"
 	"github.com/feast-dev/feast/go/internal/feast/model"
+	"github.com/feast-dev/feast/go/internal/feast/onlineserving"
 	"github.com/feast-dev/feast/go/internal/feast/server/logging"
 	"github.com/feast-dev/feast/go/protos/feast/serving"
 	prototypes "github.com/feast-dev/feast/go/protos/feast/types"
 	"github.com/feast-dev/feast/go/types"
-	"net/http"
-	"time"
 )
 
 type httpServer struct {
@@ -210,15 +212,15 @@ func (s *httpServer) getOnlineFeatures(w http.ResponseWriter, r *http.Request) {
 		"results": results,
 	}
 
+	w.Header().Set("Content-Type", "application/json")
+
 	err = json.NewEncoder(w).Encode(response)
 
 	if err != nil {
 		http.Error(w, fmt.Sprintf("Error encoding response: %+v", err), http.StatusInternalServerError)
 		return
 	}
 
-	w.Header().Set("Content-Type", "application/json")
-
 	if featureService != nil && featureService.LoggingConfig != nil && s.loggingService != nil {
 		logger, err := s.loggingService.GetOrCreateLogger(featureService)
 		if err != nil {
@@ -250,18 +252,32 @@ func (s *httpServer) getOnlineFeatures(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 	}
+	go releaseCGOMemory(featureVectors)
+}
+
+func releaseCGOMemory(featureVectors []*onlineserving.FeatureVector) {
+	for _, vector := range featureVectors {
+		vector.Values.Release()
+	}
 }
 
 func (s *httpServer) Serve(host string, port int) error {
 	s.server = &http.Server{Addr: fmt.Sprintf("%s:%d", host, port), Handler: nil}
 	http.HandleFunc("/get-online-features", s.getOnlineFeatures)
+	http.HandleFunc("/health", healthCheckHandler)
 	err := s.server.ListenAndServe()
 	// Don't return the error if it's caused by graceful shutdown using Stop()
 	if err == http.ErrServerClosed {
 		return nil
 	}
 	return err
 }
+
+func healthCheckHandler(w http.ResponseWriter, r *http.Request) {
+	w.WriteHeader(http.StatusOK)
+	fmt.Fprintf(w, "Healthy")
+}
+
 func (s *httpServer) Stop() error {
 	if s.server != nil {
 		return s.server.Shutdown(context.Background())
diff --git a/go/internal/feast/transformation/transformation.go b/go/internal/feast/transformation/transformation.go
@@ -46,13 +46,15 @@ func AugmentResponseWithOnDemandTransforms(
 		for name, values := range requestData {
 			requestContextArrow[name], err = types.ProtoValuesToArrowArray(values.Val, arrowMemory, numRows)
 			if err != nil {
+				ReleaseArrowContext(requestContextArrow)
 				return nil, err
 			}
 		}
 
 		for name, values := range entityRows {
 			requestContextArrow[name], err = types.ProtoValuesToArrowArray(values.Val, arrowMemory, numRows)
 			if err != nil {
+				ReleaseArrowContext(requestContextArrow)
 				return nil, err
 			}
 		}
@@ -71,14 +73,24 @@ func AugmentResponseWithOnDemandTransforms(
 			fullFeatureNames,
 		)
 		if err != nil {
+			ReleaseArrowContext(requestContextArrow)
 			return nil, err
 		}
 		result = append(result, onDemandFeatures...)
+
+		ReleaseArrowContext(requestContextArrow)
 	}
 
 	return result, nil
 }
 
+func ReleaseArrowContext(requestContextArrow map[string]arrow.Array) {
+	// Release memory used by requestContextArrow
+	for _, arrowArray := range requestContextArrow {
+		arrowArray.Release()
+	}
+}
+
 func CallTransformations(
 	featureView *model.OnDemandFeatureView,
 	retrievedFeatures map[string]arrow.Array,

Original file line number	Diff line number	Diff line change
`@@ -46,13 +46,15 @@ func AugmentResponseWithOnDemandTransforms(`
`46`	`46`	`for name, values := range requestData {`
`47`	`47`	`requestContextArrow[name], err = types.ProtoValuesToArrowArray(values.Val, arrowMemory, numRows)`
`48`	`48`	`if err != nil {`
	`49`	`+ ReleaseArrowContext(requestContextArrow)`
`49`	`50`	`return nil, err`
`50`	`51`	`}`
`51`	`52`	`}`
`52`	`53`
`53`	`54`	`for name, values := range entityRows {`
`54`	`55`	`requestContextArrow[name], err = types.ProtoValuesToArrowArray(values.Val, arrowMemory, numRows)`
`55`	`56`	`if err != nil {`
	`57`	`+ ReleaseArrowContext(requestContextArrow)`
`56`	`58`	`return nil, err`
`57`	`59`	`}`
`58`	`60`	`}`
`@@ -71,14 +73,24 @@ func AugmentResponseWithOnDemandTransforms(`
`71`	`73`	`fullFeatureNames,`
`72`	`74`	`)`
`73`	`75`	`if err != nil {`
	`76`	`+ ReleaseArrowContext(requestContextArrow)`
`74`	`77`	`return nil, err`
`75`	`78`	`}`
`76`	`79`	`result = append(result, onDemandFeatures...)`
	`80`	`+`
	`81`	`+ ReleaseArrowContext(requestContextArrow)`
`77`	`82`	`}`
`78`	`83`
`79`	`84`	`return result, nil`
`80`	`85`	`}`
`81`	`86`
	`87`	`+func ReleaseArrowContext(requestContextArrow map[string]arrow.Array) {`
	`88`	`+ // Release memory used by requestContextArrow`
	`89`	`+ for _, arrowArray := range requestContextArrow {`
	`90`	`+ arrowArray.Release()`
	`91`	`+ }`
	`92`	`+}`
	`93`	`+`
`82`	`94`	`func CallTransformations(`
`83`	`95`	`featureView *model.OnDemandFeatureView,`
`84`	`96`	`retrievedFeatures map[string]arrow.Array,`