🐛 fix memory leak caused by gin context pooling. (konveyor#710)

The gin HTTP request handler uses a `gin.Context` pool. The _RichContext_ contains the response body which is passed to the `Render` middle-ware. We need to remove the rich context to prevent holding on to memory while the context is in the pool. _WithContext_ renamed to _RichContext_ for clarity. Signed-off-by: Jeff Ortel <[email protected]>
aufi · Jul 8, 2024 · fac959d · fac959d
1 parent 8bf0223
commit fac959d
Show file tree

Hide file tree

Showing 11 changed files with 52 additions and 39 deletions.
diff --git a/api/analysis.go b/api/analysis.go
@@ -2403,7 +2403,7 @@ func (r *IssueWriter) Create(id uint, filter qf.Filter) (path string, count int6
 
 // db returns a db client.
 func (r *IssueWriter) db() (db *gorm.DB) {
-	rtx := WithContext(r.ctx)
+	rtx := RichContext(r.ctx)
 	db = rtx.DB.Debug()
 	return
 }
@@ -2477,7 +2477,7 @@ type AnalysisWriter struct {
 
 // db returns a db client.
 func (r *AnalysisWriter) db() (db *gorm.DB) {
-	rtx := WithContext(r.ctx)
+	rtx := RichContext(r.ctx)
 	db = rtx.DB.Debug()
 	return
 }
@@ -2615,7 +2615,7 @@ type ReportWriter struct {
 
 // db returns a db client.
 func (r *ReportWriter) db() (db *gorm.DB) {
-	rtx := WithContext(r.ctx)
+	rtx := RichContext(r.ctx)
 	db = rtx.DB.Debug()
 	return
 }

diff --git a/api/application.go b/api/application.go
@@ -249,7 +249,7 @@ func (h ApplicationHandler) Create(ctx *gin.Context) {
 		return
 	}
 
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	tr := trigger.Application{
 		Trigger: trigger.Trigger{
 			TaskManager: rtx.TaskManager,
@@ -388,7 +388,7 @@ func (h ApplicationHandler) Update(ctx *gin.Context) {
 		}
 	}
 
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	tr := trigger.Application{
 		Trigger: trigger.Trigger{
 			TaskManager: rtx.TaskManager,

diff --git a/api/auth.go b/api/auth.go
@@ -98,7 +98,7 @@ type Login struct {
 // been granted the necessary scope to access a resource.
 func Required(scope string) func(*gin.Context) {
 	return func(ctx *gin.Context) {
-		rtx := WithContext(ctx)
+		rtx := RichContext(ctx)
 		token := ctx.GetHeader(Authorization)
 		request := &auth.Request{
 			Token:  token,

diff --git a/api/base.go b/api/base.go
@@ -35,14 +35,14 @@ type BaseHandler struct{}
 
 // DB return db client associated with the context.
 func (h *BaseHandler) DB(ctx *gin.Context) (db *gorm.DB) {
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	db = rtx.DB.Debug()
 	return
 }
 
 // Client returns k8s client from the context.
 func (h *BaseHandler) Client(ctx *gin.Context) (client client.Client) {
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	client = rtx.Client
 	return
 }
@@ -100,7 +100,7 @@ func (h *BaseHandler) pk(ctx *gin.Context) (id uint) {
 
 // CurrentUser gets username from Keycloak auth token.
 func (h *BaseHandler) CurrentUser(ctx *gin.Context) (user string) {
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	user = rtx.User
 	if user == "" {
 		Log.Info("Failed to get current user.")
@@ -113,7 +113,7 @@ func (h *BaseHandler) CurrentUser(ctx *gin.Context) (user string) {
 func (h *BaseHandler) HasScope(ctx *gin.Context, scope string) (b bool) {
 	in := auth.BaseScope{}
 	in.With(scope)
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	for _, s := range rtx.Scopes {
 		b = s.Match(in.Resource, in.Method)
 		if b {
@@ -215,13 +215,13 @@ func (h *BaseHandler) Decoder(ctx *gin.Context, encoding string, r io.Reader) (d
 
 // Status sets the status code.
 func (h *BaseHandler) Status(ctx *gin.Context, code int) {
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	rtx.Status(code)
 }
 
 // Respond sets the response.
 func (h *BaseHandler) Respond(ctx *gin.Context, code int, r any) {
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	rtx.Respond(code, r)
 }
 

diff --git a/api/batch.go b/api/batch.go
@@ -62,7 +62,7 @@ func (h BatchHandler) create(ctx *gin.Context, create gin.HandlerFunc) {
 		return
 	}
 
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	bErr := BatchError{Message: "Create failed."}
 	for i := range resources {
 		b, _ := json.Marshal(resources[i])

diff --git a/api/context.go b/api/context.go
@@ -10,6 +10,12 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
+// Response values.
+type Response struct {
+	Status int
+	Body   any
+}
+
 // Context custom settings.
 type Context struct {
 	*gin.Context
@@ -27,10 +33,15 @@ type Context struct {
 	TaskManager *tasking.Manager
 }
 
-// Response values.
-type Response struct {
-	Status int
-	Body   any
+// Attach to gin context.
+func (r *Context) Attach(ctx *gin.Context) {
+	r.Context = ctx
+	ctx.Set("RichContext", r)
+}
+
+// Detach from gin context
+func (r *Context) Detach() {
+	delete(r.Context.Keys, "RichContext")
 }
 
 // Status sets the values to respond to the request with.
@@ -42,24 +53,24 @@ func (r *Context) Status(status int) {
 }
 
 // Respond sets the values to respond to the request with.
-func (r *Context) Respond(status int, body any) {
+func (r *Context) Respond(status int, body interface{}) {
 	r.Response = Response{
 		Status: status,
 		Body:   body,
 	}
 }
 
-// WithContext is a rich context.
-func WithContext(ctx *gin.Context) (n *Context) {
+// RichContext returns a rich context attached to the gin context.
+func RichContext(ctx *gin.Context) (rtx *Context) {
 	key := "RichContext"
 	object, found := ctx.Get(key)
 	if !found {
-		n = &Context{}
-		ctx.Set(key, n)
+		rtx = &Context{}
+		rtx.Attach(ctx)
 	} else {
-		n = object.(*Context)
+		rtx = object.(*Context)
 	}
-	n.Context = ctx
+	rtx.Context = ctx
 	return
 }
 
@@ -70,7 +81,7 @@ func Transaction(ctx *gin.Context) {
 		http.MethodPut,
 		http.MethodPatch,
 		http.MethodDelete:
-		rtx := WithContext(ctx)
+		rtx := RichContext(ctx)
 		err := rtx.DB.Transaction(func(tx *gorm.DB) (err error) {
 			db := rtx.DB
 			rtx.DB = tx
@@ -93,7 +104,7 @@ func Transaction(ctx *gin.Context) {
 func Render() gin.HandlerFunc {
 	return func(ctx *gin.Context) {
 		ctx.Next()
-		rtx := WithContext(ctx)
+		rtx := RichContext(ctx)
 		if rtx.Response.Body != nil {
 			ctx.Negotiate(
 				rtx.Response.Status,

diff --git a/api/error.go b/api/error.go
@@ -89,7 +89,7 @@ func ErrorHandler() gin.HandlerFunc {
 
 		err := ctx.Errors[0]
 
-		rtx := WithContext(ctx)
+		rtx := RichContext(ctx)
 		if errors.Is(err, &BadRequestError{}) ||
 			errors.Is(err, &filter.Error{}) ||
 			errors.Is(err, &sort.SortError{}) ||

diff --git a/api/identity.go b/api/identity.go
@@ -208,7 +208,7 @@ func (h IdentityHandler) Update(ctx *gin.Context) {
 		return
 	}
 
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	tr := trigger.Identity{
 		Trigger: trigger.Trigger{
 			TaskManager: rtx.TaskManager,

diff --git a/api/task.go b/api/task.go
@@ -288,7 +288,7 @@ func (h TaskHandler) Create(ctx *gin.Context) {
 		_ = ctx.Error(err)
 		return
 	}
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	task := &tasking.Task{}
 	task.With(r.Model())
 	task.CreateUser = h.BaseHandler.CurrentUser(ctx)
@@ -312,7 +312,7 @@ func (h TaskHandler) Create(ctx *gin.Context) {
 // @param id path int true "Task ID"
 func (h TaskHandler) Delete(ctx *gin.Context) {
 	id := h.pk(ctx)
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	err := rtx.TaskManager.Delete(h.DB(ctx), id)
 	if err != nil {
 		_ = ctx.Error(err)
@@ -355,7 +355,7 @@ func (h TaskHandler) Update(ctx *gin.Context) {
 	m = r.Model()
 	m.ID = id
 	m.UpdateUser = h.CurrentUser(ctx)
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	task := &tasking.Task{}
 	task.With(m)
 	err = rtx.TaskManager.Update(h.DB(ctx), task)
@@ -393,7 +393,7 @@ func (h TaskHandler) Submit(ctx *gin.Context) {
 // @param id path int true "Task ID"
 func (h TaskHandler) Cancel(ctx *gin.Context) {
 	id := h.pk(ctx)
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	err := rtx.TaskManager.Cancel(h.DB(ctx), id)
 	if err != nil {
 		_ = ctx.Error(err)

diff --git a/api/taskgroup.go b/api/taskgroup.go
@@ -142,7 +142,7 @@ func (h TaskGroupHandler) Create(ctx *gin.Context) {
 			_ = ctx.Error(result.Error)
 			return
 		}
-		rtx := WithContext(ctx)
+		rtx := RichContext(ctx)
 		for i := range m.Tasks {
 			task := &tasking.Task{}
 			task.With(&m.Tasks[i])
@@ -234,7 +234,7 @@ func (h TaskGroupHandler) Update(ctx *gin.Context) {
 			_ = ctx.Error(err)
 			return
 		}
-		rtx := WithContext(ctx)
+		rtx := RichContext(ctx)
 		for i := range m.Tasks {
 			task := &tasking.Task{}
 			task.With(&m.Tasks[i])
@@ -273,7 +273,7 @@ func (h TaskGroupHandler) Delete(ctx *gin.Context) {
 		_ = ctx.Error(err)
 		return
 	}
-	rtx := WithContext(ctx)
+	rtx := RichContext(ctx)
 	for i := range m.Tasks {
 		task := &m.Tasks[i]
 		err = rtx.TaskManager.Delete(h.DB(ctx), task.ID)

diff --git a/cmd/main.go b/cmd/main.go
@@ -198,15 +198,17 @@ func main() {
 	}
 	// Web
 	router := gin.Default()
-	router.Use(api.Render())
-	router.Use(api.ErrorHandler())
 	router.Use(
 		func(ctx *gin.Context) {
-			rtx := api.WithContext(ctx)
+			rtx := api.RichContext(ctx)
+			rtx.TaskManager = &taskManager
 			rtx.DB = db
 			rtx.Client = client
-			rtx.TaskManager = &taskManager
+			ctx.Next()
+			rtx.Detach()
 		})
+	router.Use(api.Render())
+	router.Use(api.ErrorHandler())
 	for _, h := range api.All() {
 		h.AddRoutes(router)
 	}