From f75a90d88d348e45b7019ef5f1a462046a578ec9 Mon Sep 17 00:00:00 2001
From: chrislu <chris.lu@gmail.com>
Date: Thu, 20 Nov 2025 16:47:02 -0800
Subject: [PATCH] fix: goroutine and connection leak in IAM server shutdown

The IAM server's KeepConnectedToMaster goroutine used context.Background(),
which is non-cancellable, causing the goroutine and its gRPC connections
to leak on server shutdown.

Problem:
  go masterClient.KeepConnectedToMaster(context.Background())

  - context.Background() never cancels
  - KeepConnectedToMaster goroutine runs forever
  - gRPC connection to master stays open
  - No way to stop cleanly on server shutdown

Result: Resource leaks when IAM server is stopped

Fix:
  1. Added shutdownContext and shutdownCancel to IamApiServer struct
  2. Created cancellable context in NewIamApiServerWithStore:
       shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
  3. Pass shutdownCtx to KeepConnectedToMaster:
       go masterClient.KeepConnectedToMaster(shutdownCtx)
  4. Added Shutdown() method to invoke cancel:
       func (iama *IamApiServer) Shutdown() {
           if iama.shutdownCancel != nil {
               iama.shutdownCancel()
           }
       }

  5. Stored masterClient reference on IamApiServer for future use

Benefits:
  - Goroutine stops cleanly when Shutdown() is called
  - gRPC connections are closed properly
  - No resource leaks on server restart/stop
  - Shutdown() is idempotent (safe to call multiple times)

Usage (for future graceful shutdown):
  iamServer, _ := iamapi.NewIamApiServer(...)
  defer iamServer.Shutdown()

  // or in signal handler:
  sigChan := make(chan os.Signal, 1)
  signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
  go func() {
      <-sigChan
      iamServer.Shutdown()
      os.Exit(0)
  }()

Note: Current command implementations (weed/command/iam.go) don't have
shutdown paths yet, but this makes IAM server ready for proper lifecycle
management when that infrastructure is added.
---
 weed/iamapi/iamapi_server.go | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/weed/iamapi/iamapi_server.go b/weed/iamapi/iamapi_server.go
index 82977458d..94dd599ba 100644
--- a/weed/iamapi/iamapi_server.go
+++ b/weed/iamapi/iamapi_server.go
@@ -45,8 +45,11 @@ type IamServerOption struct {
 }
 
 type IamApiServer struct {
-	s3ApiConfig IamS3ApiConfig
-	iam         *s3api.IdentityAccessManagement
+	s3ApiConfig      IamS3ApiConfig
+	iam              *s3api.IdentityAccessManagement
+	shutdownContext  context.Context
+	shutdownCancel   context.CancelFunc
+	masterClient     *wdclient.MasterClient
 }
 
 var s3ApiConfigure IamS3ApiConfig
@@ -58,10 +61,14 @@ func NewIamApiServer(router *mux.Router, option *IamServerOption) (iamApiServer
 func NewIamApiServerWithStore(router *mux.Router, option *IamServerOption, explicitStore string) (iamApiServer *IamApiServer, err error) {
 	masterClient := wdclient.NewMasterClient(option.GrpcDialOption, "", "iam", "", "", "", *pb.NewServiceDiscoveryFromMap(option.Masters))
 	
+	// Create a cancellable context for the master client connection
+	// This allows graceful shutdown via Shutdown() method
+	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
+	
 	// Start KeepConnectedToMaster for volume location lookups
 	// IAM config files are typically small and inline, but if they ever have chunks,
 	// ReadEntry→StreamContent needs masterClient for volume lookups
-	go masterClient.KeepConnectedToMaster(context.Background())
+	go masterClient.KeepConnectedToMaster(shutdownCtx)
 	
 	configure := &IamS3ApiConfigure{
 		option:       option,
@@ -79,8 +86,11 @@ func NewIamApiServerWithStore(router *mux.Router, option *IamServerOption, expli
 	configure.credentialManager = iam.GetCredentialManager()
 
 	iamApiServer = &IamApiServer{
-		s3ApiConfig: s3ApiConfigure,
-		iam:         iam,
+		s3ApiConfig:     s3ApiConfigure,
+		iam:             iam,
+		shutdownContext: shutdownCtx,
+		shutdownCancel:  shutdownCancel,
+		masterClient:    masterClient,
 	}
 
 	iamApiServer.registerRouter(router)
@@ -100,6 +110,15 @@ func (iama *IamApiServer) registerRouter(router *mux.Router) {
 	apiRouter.NotFoundHandler = http.HandlerFunc(s3err.NotFoundHandler)
 }
 
+// Shutdown gracefully stops the IAM API server and releases resources.
+// It cancels the master client connection goroutine and closes gRPC connections.
+// This method is safe to call multiple times.
+func (iama *IamApiServer) Shutdown() {
+	if iama.shutdownCancel != nil {
+		iama.shutdownCancel()
+	}
+}
+
 func (iama *IamS3ApiConfigure) GetS3ApiConfiguration(s3cfg *iam_pb.S3ApiConfiguration) (err error) {
 	return iama.GetS3ApiConfigurationFromCredentialManager(s3cfg)
 }