Skip to content

Commit 43e198f

Browse files
authored
fix: CGO Memory leak issue in GO Feature server (#4291)
1 parent 2c38946 commit 43e198f

File tree

2 files changed

+32
-4
lines changed

2 files changed

+32
-4
lines changed

go/internal/feast/server/http_server.go

+20-4
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,16 @@ import (
44
"context"
55
"encoding/json"
66
"fmt"
7+
"net/http"
8+
"time"
9+
710
"github.com/feast-dev/feast/go/internal/feast"
811
"github.com/feast-dev/feast/go/internal/feast/model"
12+
"github.com/feast-dev/feast/go/internal/feast/onlineserving"
913
"github.com/feast-dev/feast/go/internal/feast/server/logging"
1014
"github.com/feast-dev/feast/go/protos/feast/serving"
1115
prototypes "github.com/feast-dev/feast/go/protos/feast/types"
1216
"github.com/feast-dev/feast/go/types"
13-
"net/http"
14-
"time"
1517
)
1618

1719
type httpServer struct {
@@ -210,15 +212,15 @@ func (s *httpServer) getOnlineFeatures(w http.ResponseWriter, r *http.Request) {
210212
"results": results,
211213
}
212214

215+
w.Header().Set("Content-Type", "application/json")
216+
213217
err = json.NewEncoder(w).Encode(response)
214218

215219
if err != nil {
216220
http.Error(w, fmt.Sprintf("Error encoding response: %+v", err), http.StatusInternalServerError)
217221
return
218222
}
219223

220-
w.Header().Set("Content-Type", "application/json")
221-
222224
if featureService != nil && featureService.LoggingConfig != nil && s.loggingService != nil {
223225
logger, err := s.loggingService.GetOrCreateLogger(featureService)
224226
if err != nil {
@@ -250,18 +252,32 @@ func (s *httpServer) getOnlineFeatures(w http.ResponseWriter, r *http.Request) {
250252
return
251253
}
252254
}
255+
go releaseCGOMemory(featureVectors)
256+
}
257+
258+
func releaseCGOMemory(featureVectors []*onlineserving.FeatureVector) {
259+
for _, vector := range featureVectors {
260+
vector.Values.Release()
261+
}
253262
}
254263

255264
func (s *httpServer) Serve(host string, port int) error {
256265
s.server = &http.Server{Addr: fmt.Sprintf("%s:%d", host, port), Handler: nil}
257266
http.HandleFunc("/get-online-features", s.getOnlineFeatures)
267+
http.HandleFunc("/health", healthCheckHandler)
258268
err := s.server.ListenAndServe()
259269
// Don't return the error if it's caused by graceful shutdown using Stop()
260270
if err == http.ErrServerClosed {
261271
return nil
262272
}
263273
return err
264274
}
275+
276+
func healthCheckHandler(w http.ResponseWriter, r *http.Request) {
277+
w.WriteHeader(http.StatusOK)
278+
fmt.Fprintf(w, "Healthy")
279+
}
280+
265281
func (s *httpServer) Stop() error {
266282
if s.server != nil {
267283
return s.server.Shutdown(context.Background())

go/internal/feast/transformation/transformation.go

+12
Original file line numberDiff line numberDiff line change
@@ -46,13 +46,15 @@ func AugmentResponseWithOnDemandTransforms(
4646
for name, values := range requestData {
4747
requestContextArrow[name], err = types.ProtoValuesToArrowArray(values.Val, arrowMemory, numRows)
4848
if err != nil {
49+
ReleaseArrowContext(requestContextArrow)
4950
return nil, err
5051
}
5152
}
5253

5354
for name, values := range entityRows {
5455
requestContextArrow[name], err = types.ProtoValuesToArrowArray(values.Val, arrowMemory, numRows)
5556
if err != nil {
57+
ReleaseArrowContext(requestContextArrow)
5658
return nil, err
5759
}
5860
}
@@ -71,14 +73,24 @@ func AugmentResponseWithOnDemandTransforms(
7173
fullFeatureNames,
7274
)
7375
if err != nil {
76+
ReleaseArrowContext(requestContextArrow)
7477
return nil, err
7578
}
7679
result = append(result, onDemandFeatures...)
80+
81+
ReleaseArrowContext(requestContextArrow)
7782
}
7883

7984
return result, nil
8085
}
8186

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+
8294
func CallTransformations(
8395
featureView *model.OnDemandFeatureView,
8496
retrievedFeatures map[string]arrow.Array,

0 commit comments

Comments
 (0)