More resilient droplet upload (#195)

* More resilient droplet upload, ensure draining for complete upload process

Author: kathap

ccclient/poller.go

@@ -47,6 +47,7 @@ func (p *poller) Poll(fallbackURL *url.URL, res *http.Response, cancelChan <-cha
 
 		switch body.Entity.Status {
 		case JOB_QUEUED, JOB_RUNNING:
+			p.logger.Info("cc-job-queued-or-running", lager.Data{"status": body.Entity.Status})
 		case JOB_FINISHED:
 			p.logger.Info("cc-job-finished")
 			return nil

cmd/cc-uploader/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"code.cloudfoundry.org/tlsconfig"
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"flag"
@@ -11,7 +12,10 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"os/signal"
 	"runtime"
+	"sync"
+	"syscall"
 	"time"
 
 	"code.cloudfoundry.org/cc-uploader/ccclient"
@@ -41,6 +45,9 @@ const (
 	communicationTimeout        = 30 * time.Second
 )
 
+// Global WaitGroup to track uploads
+var uploadWaitGroup sync.WaitGroup
+
 func main() {
 	runtime.GOMAXPROCS(runtime.NumCPU())
 	flag.Parse()
@@ -54,12 +61,20 @@ func main() {
 
 	initializeDropsonde(logger, uploaderConfig)
 
+	// Create signal channel to listen for shutdown signals
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, os.Interrupt, syscall.SIGTERM)
+
+	var nonTLSServer *http.Server
+	tlsServer, tlsRunner := initializeServer(logger, uploaderConfig, true)
 	members := grouper.Members{
-		{"cc-uploader-tls", initializeServer(logger, uploaderConfig, true)},
+		{"cc-uploader-tls", tlsRunner},
 	}
 	if !uploaderConfig.DisableNonTLS {
+		var nonTLSRunner ifrit.Runner
+		nonTLSServer, nonTLSRunner = initializeServer(logger, uploaderConfig, false)
 		members = append(grouper.Members{
-			{"cc-uploader", initializeServer(logger, uploaderConfig, false)},
+			{"cc-uploader", nonTLSRunner},
 		}, members...)
 	}
 	if uploaderConfig.DebugServerConfig.DebugAddress != "" {
@@ -73,10 +88,32 @@ func main() {
 	monitor := ifrit.Invoke(sigmon.New(group))
 	logger.Info("ready")
 
-	err = <-monitor.Wait()
-	if err != nil {
-		logger.Error("exited-with-failure", err)
-		os.Exit(1)
+	select {
+	case err := <-monitor.Wait():
+		if err != nil {
+			logger.Info("exited-with-failure")
+			os.Exit(1)
+		}
+	case sig := <-signalChan:
+		logger.Info("shutdown-signal-received", lager.Data{"signal": sig})
+
+		// Gracefully signal Ifrit monitor to stop processes
+		monitor.Signal(os.Interrupt)
+		logger.Info("graceful-shutdown-waiting-for-uploads")
+		// Wait for all uploads to finish before shutting down
+		uploadWaitGroup.Wait()
+		// Gracefully shutdown the HTTP server
+		ctx, cancel := context.WithTimeout(context.Background(), 300*time.Second)
+		defer cancel()
+		if !uploaderConfig.DisableNonTLS {
+			if err := nonTLSServer.Shutdown(ctx); err != nil {
+				logger.Error("non-tls-server-shutdown-failed", err)
+			}
+		}
+
+		if err := tlsServer.Shutdown(ctx); err != nil {
+			logger.Error("tls-server-shutdown-failed", err)
+		}
 	}
 
 	logger.Info("exited")
@@ -123,18 +160,19 @@ func initializeTlsTransport(uploaderConfig config.UploaderConfig, skipVerify boo
 	}
 }
 
-func initializeServer(logger lager.Logger, uploaderConfig config.UploaderConfig, tlsServer bool) ifrit.Runner {
+func initializeServer(logger lager.Logger, uploaderConfig config.UploaderConfig, tlsServer bool) (*http.Server, ifrit.Runner) {
 	uploader := ccclient.NewUploader(logger, &http.Client{Transport: initializeTlsTransport(uploaderConfig, false)})
 
 	// To maintain backwards compatibility with hairpin polling URLs, skip SSL verification for now
 	poller := ccclient.NewPoller(logger, &http.Client{Transport: initializeTlsTransport(uploaderConfig, true)}, time.Duration(uploaderConfig.CCJobPollingInterval))
 
-	ccUploaderHandler, err := handlers.New(uploader, poller, logger)
+	ccUploaderHandler, err := handlers.New(uploader, poller, logger, &uploadWaitGroup)
 	if err != nil {
 		logger.Error("router-building-failed", err)
 		os.Exit(1)
 	}
 
+	var server *http.Server
 	if tlsServer {
 		clientTLSConfig, err := tlsconfig.Build(
 			tlsconfig.WithIdentityFromFile(uploaderConfig.MutualTLS.ServerCert, uploaderConfig.MutualTLS.ServerKey),
@@ -151,10 +189,16 @@ func initializeServer(logger lager.Logger, uploaderConfig config.UploaderConfig,
 			tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
 		}
 
-		if err != nil {
-			logger.Fatal("failed-loading-tls-config", err)
+		server = &http.Server{
+			Addr:      uploaderConfig.MutualTLS.ListenAddress,
+			Handler:   ccUploaderHandler,
+			TLSConfig: clientTLSConfig,
 		}
-		return http_server.NewTLSServer(uploaderConfig.MutualTLS.ListenAddress, ccUploaderHandler, clientTLSConfig)
+		return server, http_server.NewTLSServer(uploaderConfig.MutualTLS.ListenAddress, ccUploaderHandler, clientTLSConfig)
+	}
+	server = &http.Server{
+		Addr:    uploaderConfig.ListenAddress,
+		Handler: ccUploaderHandler,
 	}
-	return http_server.New(uploaderConfig.ListenAddress, ccUploaderHandler)
+	return server, http_server.New(uploaderConfig.ListenAddress, ccUploaderHandler)
 }

cmd/cc-uploader/main_test.go

@@ -257,6 +257,34 @@ var _ = Describe("CC Uploader", func() {
 					Expect(len(fakeCC.UploadedDroplets[appGuid])).To(Equal(contentLength))
 				})
 			})
+
 		})
 	})
+
+	Describe("Handling shutdown signals", func() {
+		It("should handle SIGTERM and drain ongoing uploads before shutting down", func() {
+			ccUploaderAddress := fmt.Sprintf("http://localhost:%d", httpListenPort)
+			emitter := NewEmitter(100) // large, slow upload
+			postRequest := dropletUploadRequest(appGuid, emitter, 100, ccUploaderAddress)
+
+			go func() {
+				_, err := http.DefaultClient.Do(postRequest)
+				Expect(err).NotTo(HaveOccurred())
+			}()
+
+			// Give it a moment to start the upload
+			time.Sleep(200 * time.Millisecond)
+
+			// Send SIGTERM to the cc-uploader process
+			session.Signal(os.Interrupt)
+
+			// Expect shutdown logs
+			Eventually(session, 1*time.Second).Should(gbytes.Say("shutdown-signal-received"))
+			Eventually(session, 1*time.Second).Should(gbytes.Say("graceful-shutdown-waiting-for-uploads"))
+
+			// Wait for the process to exit cleanly
+			Eventually(session, 2*time.Second).Should(gexec.Exit(0))
+		})
+	})
+
 })

handlers/handlers.go

@@ -2,6 +2,7 @@ package handlers
 
 import (
 	"net/http"
+	"sync"
 
 	"code.cloudfoundry.org/cc-uploader"
 	"code.cloudfoundry.org/cc-uploader/ccclient"
@@ -11,9 +12,9 @@ import (
 	"github.com/tedsuo/rata"
 )
 
-func New(uploader ccclient.Uploader, poller ccclient.Poller, logger lager.Logger) (http.Handler, error) {
+func New(uploader ccclient.Uploader, poller ccclient.Poller, logger lager.Logger, uploadWaitGroup *sync.WaitGroup) (http.Handler, error) {
 	return rata.NewRouter(ccuploader.Routes, rata.Handlers{
-		ccuploader.UploadDropletRoute:        upload_droplet.New(uploader, poller, logger),
+		ccuploader.UploadDropletRoute:        upload_droplet.New(uploader, poller, logger, uploadWaitGroup),
 		ccuploader.UploadBuildArtifactsRoute: upload_build_artifacts.New(uploader, logger),
 	})
 }

handlers/handlers_test.go

@@ -8,6 +8,7 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"strconv"
+	"sync"
 	"time"
 
 	"code.cloudfoundry.org/cc-uploader/ccclient"
@@ -53,8 +54,8 @@ var _ = Describe("Handlers", func() {
 
 		uploader := ccclient.NewUploader(logger, http.DefaultClient)
 		poller := ccclient.NewPoller(logger, http.DefaultClient, 100*time.Millisecond)
-
-		handler, err = handlers.New(uploader, poller, logger)
+		var wg sync.WaitGroup
+		handler, err = handlers.New(uploader, poller, logger, &wg)
 		Expect(err).NotTo(HaveOccurred())
 
 		postStatusCode = http.StatusCreated

handlers/upload_droplet/upload_droplet.go

@@ -6,32 +6,37 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"sync"
 	"time"
 
 	"code.cloudfoundry.org/cc-uploader/ccclient"
 	"code.cloudfoundry.org/lager/v3"
 	"code.cloudfoundry.org/runtimeschema/cc_messages"
 )
 
-func New(uploader ccclient.Uploader, poller ccclient.Poller, logger lager.Logger) http.Handler {
+func New(uploader ccclient.Uploader, poller ccclient.Poller, logger lager.Logger, uploadWaitGroup *sync.WaitGroup) http.Handler {
 	return &dropletUploader{
-		uploader: uploader,
-		poller:   poller,
-		logger:   logger,
+		uploader:        uploader,
+		poller:          poller,
+		logger:          logger,
+		uploadWaitGroup: uploadWaitGroup,
 	}
 }
 
 type dropletUploader struct {
-	uploader ccclient.Uploader
-	poller   ccclient.Poller
-	logger   lager.Logger
+	uploader        ccclient.Uploader
+	poller          ccclient.Poller
+	logger          lager.Logger
+	uploadWaitGroup *sync.WaitGroup
 }
 
 var MissingCCDropletUploadUriKeyError = errors.New(fmt.Sprintf("missing %s parameter", cc_messages.CcDropletUploadUriKey))
 
 func (h *dropletUploader) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	logger := h.logger.Session("droplet.upload")
-
+	h.uploadWaitGroup.Add(1)
+	// Ensure that the WaitGroup is decremented when the function returns
+	defer h.uploadWaitGroup.Done()
 	logger.Info("extracting-droplet-upload-uri-key")
 	uploadUriParameter := r.URL.Query().Get(cc_messages.CcDropletUploadUriKey)
 	if uploadUriParameter == "" {

handlers/upload_droplet/upload_droplet_test.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"sync"
 	"time"
 
 	"code.cloudfoundry.org/cc-uploader/ccclient/fake_ccclient"
@@ -36,7 +37,8 @@ var _ = Describe("UploadDroplet", func() {
 
 		JustBeforeEach(func() {
 			logger = lager.NewLogger("fake-logger")
-			dropletUploadHandler := upload_droplet.New(&uploader, &poller, logger)
+			var wg sync.WaitGroup
+			dropletUploadHandler := upload_droplet.New(&uploader, &poller, logger, &wg)
 
 			dropletUploadHandler.ServeHTTP(responseWriter, incomingRequest)
 		})