Improve droplet upload draining

ccclient/uploader.go

@@ -53,7 +53,7 @@ func (u *uploader) Upload(uploadURL *url.URL, filename string, r *http.Request,
 			logger.Info("succeeded-uploading")
 			break
 		}
-		logger.Error("failed-uploading", err)
+		logger.Error("failed-uploading", uploadErr)
 
 		// not a connect (dial) error
 		var nestedErr error = uploadErr

cmd/cc-uploader/main.go

@@ -1,27 +1,25 @@
 package main
 
 import (
-	"code.cloudfoundry.org/tlsconfig"
-	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"flag"
 	"fmt"
-	"io/ioutil"
 	"log"
 	"net"
 	"net/http"
 	"os"
 	"os/signal"
-	"runtime"
 	"sync"
 	"syscall"
 	"time"
 
+	"code.cloudfoundry.org/debugserver"
+	"code.cloudfoundry.org/tlsconfig"
+
 	"code.cloudfoundry.org/cc-uploader/ccclient"
 	"code.cloudfoundry.org/cc-uploader/config"
 	"code.cloudfoundry.org/cc-uploader/handlers"
-	"code.cloudfoundry.org/debugserver"
 	"code.cloudfoundry.org/lager/v3"
 	"code.cloudfoundry.org/lager/v3/lagerflags"
 	"github.com/cloudfoundry/dropsonde"
@@ -37,86 +35,29 @@ var configPath = flag.String(
 	"path to config",
 )
 
+var shutdownTimeoutInMinutes = flag.Int(
+	"shutdownTimeoutInMinutes",
+	15,
+	"max time (minutes) to wait for graceful shutdown",
+)
+
 const (
 	ccUploadDialTimeout         = 10 * time.Second
 	ccUploadKeepAlive           = 30 * time.Second
 	ccUploadTLSHandshakeTimeout = 10 * time.Second
 	dropsondeOrigin             = "cc_uploader"
-	communicationTimeout        = 30 * time.Second
 )
 
-// Global WaitGroup to track uploads
-var uploadWaitGroup sync.WaitGroup
-
-func main() {
-	runtime.GOMAXPROCS(runtime.NumCPU())
-	flag.Parse()
-
-	uploaderConfig, err := config.NewUploaderConfig(*configPath)
-	if err != nil {
-		panic(err.Error())
-	}
-
-	logger, reconfigurableSink := lagerflags.NewFromConfig("cc-uploader", uploaderConfig.LagerConfig)
-
-	initializeDropsonde(logger, uploaderConfig)
+var (
+	// Global WaitGroup to track uploads
+	uploadWaitGroup sync.WaitGroup
+)
 
+func newShutdownSignalChannel() <-chan os.Signal {
 	// 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", tlsRunner},
-	}
-	if !uploaderConfig.DisableNonTLS {
-		var nonTLSRunner ifrit.Runner
-		nonTLSServer, nonTLSRunner = initializeServer(logger, uploaderConfig, false)
-		members = append(grouper.Members{
-			{"cc-uploader", nonTLSRunner},
-		}, members...)
-	}
-	if uploaderConfig.DebugServerConfig.DebugAddress != "" {
-		members = append(grouper.Members{
-			{"debug-server", debugserver.Runner(uploaderConfig.DebugServerConfig.DebugAddress, reconfigurableSink)},
-		}, members...)
-	}
-
-	group := grouper.NewOrdered(os.Interrupt, members)
-
-	monitor := ifrit.Invoke(sigmon.New(group))
-	logger.Info("ready")
-
-	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")
+	s := make(chan os.Signal, 1)
+	signal.Notify(s, os.Interrupt, syscall.SIGTERM)
+	return s
 }
 
 func initializeDropsonde(logger lager.Logger, uploaderConfig config.UploaderConfig) {
@@ -133,7 +74,7 @@ func initializeTlsTransport(uploaderConfig config.UploaderConfig, skipVerify boo
 		log.Fatalln("Unable to load cert", err)
 	}
 
-	clientCACert, err := ioutil.ReadFile(uploaderConfig.CCCACert)
+	clientCACert, err := os.ReadFile(uploaderConfig.CCCACert)
 	if err != nil {
 		log.Fatal("Unable to open cert", err)
 	}
@@ -160,7 +101,7 @@ func initializeTlsTransport(uploaderConfig config.UploaderConfig, skipVerify boo
 	}
 }
 
-func initializeServer(logger lager.Logger, uploaderConfig config.UploaderConfig, tlsServer bool) (*http.Server, ifrit.Runner) {
+func initializeServer(logger lager.Logger, uploaderConfig config.UploaderConfig, tlsServer bool) 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
@@ -172,7 +113,6 @@ func initializeServer(logger lager.Logger, uploaderConfig config.UploaderConfig,
 		os.Exit(1)
 	}
 
-	var server *http.Server
 	if tlsServer {
 		clientTLSConfig, err := tlsconfig.Build(
 			tlsconfig.WithIdentityFromFile(uploaderConfig.MutualTLS.ServerCert, uploaderConfig.MutualTLS.ServerKey),
@@ -189,16 +129,85 @@ func initializeServer(logger lager.Logger, uploaderConfig config.UploaderConfig,
 			tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
 		}
 
-		server = &http.Server{
-			Addr:      uploaderConfig.MutualTLS.ListenAddress,
-			Handler:   ccUploaderHandler,
-			TLSConfig: clientTLSConfig,
-		}
-		return server, http_server.NewTLSServer(uploaderConfig.MutualTLS.ListenAddress, ccUploaderHandler, clientTLSConfig)
+		return http_server.NewTLSServer(uploaderConfig.MutualTLS.ListenAddress, ccUploaderHandler, clientTLSConfig)
+	}
+
+	return http_server.New(uploaderConfig.ListenAddress, ccUploaderHandler)
+}
+
+func waitForDrainingToFinish() <-chan struct{} {
+	done := make(chan struct{})
+	go func() {
+		uploadWaitGroup.Wait() // wait for all in-flight uploads to call Done()
+		close(done)            // signal completion
+	}()
+	return done
+}
+
+func configureServers(logger lager.Logger, uploaderConfig config.UploaderConfig, reconfigurableSink *lager.ReconfigurableSink) ifrit.Process {
+
+	var nonTLSRunner ifrit.Runner
+	tlsRunner := initializeServer(logger, uploaderConfig, true)
+	members := grouper.Members{
+		{Name: "cc-uploader-tls", Runner: tlsRunner},
+	}
+	if !uploaderConfig.DisableNonTLS {
+		nonTLSRunner = initializeServer(logger, uploaderConfig, false)
+		members = append(grouper.Members{
+			{Name: "cc-uploader", Runner: nonTLSRunner},
+		}, members...)
+	}
+	if uploaderConfig.DebugServerConfig.DebugAddress != "" {
+		members = append(grouper.Members{
+			{Name: "debug-server", Runner: debugserver.Runner(uploaderConfig.DebugServerConfig.DebugAddress, reconfigurableSink)},
+		}, members...)
+	}
+
+	group := grouper.NewOrdered(os.Interrupt, members)
+	monitor := ifrit.Invoke(sigmon.New(group))
+	logger.Info("ready")
+
+	return monitor
+}
+
+func main() {
+	flag.Parse()
+
+	uploaderConfig, err := config.NewUploaderConfig(*configPath)
+	if err != nil {
+		log.Fatal(err)
 	}
-	server = &http.Server{
-		Addr:    uploaderConfig.ListenAddress,
-		Handler: ccUploaderHandler,
+
+	logger, reconfigurableSink := lagerflags.NewFromConfig("cc-uploader", uploaderConfig.LagerConfig)
+
+	initializeDropsonde(logger, uploaderConfig)
+
+	shutdownSignal := newShutdownSignalChannel()
+
+	monitor := configureServers(logger, uploaderConfig, reconfigurableSink)
+
+	select {
+	case err := <-monitor.Wait():
+		if err != nil {
+			log.Fatal("server-exited-with-failure: ", err)
+		}
+	case s := <-shutdownSignal:
+		logger.Info("shutdown-signal-received", lager.Data{"signal": s})
+
+		// Stop accepting new connections on both runners (TLS & non-TLS), Ifrit will close the listeners
+		monitor.Signal(os.Interrupt)
+
+		// Create channel to signal when uploads (including polling) are done
+		done := waitForDrainingToFinish()
+
+		select {
+		case <-done:
+			logger.Info("all-uploads-finished")
+		case <-time.After(time.Duration(*shutdownTimeoutInMinutes) * time.Minute):
+			logger.Info("graceful-shutdown-timed-out",
+				lager.Data{"timeout": fmt.Sprintf("%d minutes", *shutdownTimeoutInMinutes)})
+		}
 	}
-	return server, http_server.New(uploaderConfig.ListenAddress, ccUploaderHandler)
+
+	logger.Info("exited")
 }

cmd/cc-uploader/main_test.go

@@ -58,6 +58,7 @@ var _ = Describe("CC Uploader", func() {
 		configFile      *os.File
 		appGuid         = "app-guid"
 		fakeCCServer    *httptest.Server
+		extraArgs       []string
 	)
 
 	dropletUploadRequest := func(appGuid string, body io.Reader, contentLength int, address string) *http.Request {
@@ -116,6 +117,7 @@ var _ = Describe("CC Uploader", func() {
 		args := []string{
 			"-configPath", configFile.Name(),
 		}
+		args = append(args, extraArgs...)
 		session, err = gexec.Start(exec.Command(ccUploaderBinary, args...), GinkgoWriter, GinkgoWriter)
 		Expect(err).NotTo(HaveOccurred())
 
@@ -124,6 +126,7 @@ var _ = Describe("CC Uploader", func() {
 
 	AfterEach(func() {
 		os.Remove(configFile.Name())
+		extraArgs = nil
 		session.Kill().Wait()
 	})
 
@@ -280,11 +283,85 @@ var _ = Describe("CC Uploader", func() {
 
 			// 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
+			// Now we should see that in-flight uploads actually finished
+			Eventually(session, 2*time.Second).Should(gbytes.Say("all-uploads-finished"))
+
+			// And finally the process exits cleanly
+			Eventually(session, 2*time.Second).Should(gexec.Exit(0))
+		})
+	})
+
+	Describe("Draining timeout", func() {
+		var ccUploaderAddress string
+
+		BeforeEach(func() {
+			// tiny shutdown timeout for this spec
+			extraArgs = []string{"-shutdownTimeoutInMinutes", "0"}
+
+			ccUploaderAddress = fmt.Sprintf("http://localhost:%d", httpListenPort)
+			// Only for the timeout spec
+			fakeCCServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				_, _ = io.Copy(io.Discard, r.Body) // ignore partial read errors
+				w.WriteHeader(http.StatusCreated)
+			}))
+		})
+
+		AfterEach(func() {
+			fakeCCServer.Close()
+		})
+
+		It("logs graceful-shutdown-timed-out and exits cleanly", func() {
+			// very large, very slow body: ByteEmitter emits 1 byte per ms
+			contentLength := 200_000 // ~200s if fully sent; we’ll time out long before
+			emitter := NewEmitter(contentLength)
+			req := dropletUploadRequest(appGuid, emitter, contentLength, ccUploaderAddress)
+
+			// start the upload in background
+			go func() { _, _ = http.DefaultClient.Do(req) }()
+
+			// give it a moment to enter the handler and increment the WaitGroup
+			time.Sleep(150 * time.Millisecond)
+
+			// send SIGTERM to trigger drain
+			session.Signal(os.Interrupt)
+
+			Eventually(session, time.Second).Should(gbytes.Say("shutdown-signal-received"))
+			// we expect the timeout path, not "all-uploads-finished"
+			Eventually(session, 2*time.Second).Should(gbytes.Say("graceful-shutdown-timed-out"))
 			Eventually(session, 2*time.Second).Should(gexec.Exit(0))
 		})
 	})
 
+	Describe("After SIGTERM", func() {
+		var ccUploaderAddress string
+
+		BeforeEach(func() {
+			fakeCC = fake_cc.New()
+			fakeCCServer = httptest.NewUnstartedServer(fakeCC)
+			fakeCCServer.Start()
+			ccUploaderAddress = fmt.Sprintf("http://localhost:%d", httpListenPort)
+		})
+
+		AfterEach(func() { fakeCCServer.Close() })
+
+		It("stops accepting new connections", func() {
+			// kick off one in-flight request (so process is alive)
+			emitter := NewEmitter(1000)
+			inFlight := dropletUploadRequest(appGuid, emitter, 1000, ccUploaderAddress)
+			go func() { _, _ = http.DefaultClient.Do(inFlight) }()
+
+			time.Sleep(100 * time.Millisecond)
+			session.Signal(os.Interrupt)
+			Eventually(session, time.Second).Should(gbytes.Say("shutdown-signal-received"))
+
+			// a new request should fail quickly
+			newReq := dropletUploadRequest("another-app", NewEmitter(10), 10, ccUploaderAddress)
+			// Allow either a transport error or a non-2xx quickly
+			client := &http.Client{Timeout: 500 * time.Millisecond}
+			_, err := client.Do(newReq)
+			Expect(err).To(HaveOccurred())
+		})
+	})
+
 })

handlers/handlers.go

@@ -12,7 +12,7 @@ import (
 	"github.com/tedsuo/rata"
 )
 
-func New(uploader ccclient.Uploader, poller ccclient.Poller, logger lager.Logger, uploadWaitGroup *sync.WaitGroup) (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, uploadWaitGroup),
 		ccuploader.UploadBuildArtifactsRoute: upload_build_artifacts.New(uploader, logger),