Add drain operation + endpoint to clean properly the server before shutdown

cmd/drain.go

@@ -0,0 +1,27 @@
+package cmd
+
+import (
+	"github.com/runatlantis/atlantis/drain"
+	"github.com/runatlantis/atlantis/server/logging"
+	"github.com/spf13/cobra"
+)
+
+// DrainCmd performs a drain of the local Atlantis server for all running operations.
+// The server itself is not shutdown but drained from all running operations.
+// When the command returns, the "atlantis server" process can be stopped securely.
+type DrainCmd struct {
+	Logger *logging.SimpleLogger
+}
+
+// Drain returns the runnable cobra command.
+func (v *DrainCmd) Init() *cobra.Command {
+	return &cobra.Command{
+		Use:   "drain",
+		Short: "Perform a drain of the local Atlantis server, waiting for completion before returning",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			err := drain.Start(v.Logger)
+			return err
+		},
+		SilenceErrors: true,
+	}
+}

drain/drain.go

@@ -0,0 +1,103 @@
+package drain
+
+import (
+	"encoding/json"
+	"errors"
+	"io/ioutil"
+	"net/http"
+	"time"
+
+	"github.com/runatlantis/atlantis/server"
+	"github.com/runatlantis/atlantis/server/logging"
+)
+
+// Start begins the shutdown process.
+func Start(logger *logging.SimpleLogger) error {
+	logger.Info("Drain starting")
+	http_client := &http.Client{}
+	resp, err := startDrain(http_client, logger)
+	if err != nil {
+		return err
+	}
+	logger.Info("Drain of server initiated succesfully")
+	for {
+		if resp.DrainCompleted {
+			logger.Info("Drain of server completed successfully. You can now send a TERM signal to the server.")
+			break
+		}
+		logger.Info("Drain of server still ongoing, waiting a little bit ...")
+		time.Sleep(5 * time.Second)
+		resp, err = getDrainStatus(http_client, logger)
+	}
+	return nil
+}
+
+func startDrain(http_client *http.Client, logger *logging.SimpleLogger) (*server.DrainResponse, error) {
+
+	req, err := http.NewRequest("POST", "http://localhost:4141/drain", nil)
+	if err != nil {
+		logger.Err("Failed to create POST request to /drain endpoint: %s", err)
+		return nil, err
+	}
+
+	resp, err := http_client.Do(req)
+	if err != nil {
+		logger.Err("Failed to make POST request to /drain endpoint: %s", err)
+		return nil, err
+	}
+	defer resp.Body.Close()
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		logger.Err("Failed to read reponse body of POST request to /drain endpoint: %s", err)
+		return nil, err
+	}
+
+	if resp.StatusCode != http.StatusCreated {
+		logger.Err("Unexpected status code while making POST request to /drain endpoint: ", resp.StatusCode)
+		logger.Info("Response content: %s", string(body))
+		return nil, errors.New("Unexpected status code")
+	}
+
+	var response server.DrainResponse
+	err = json.Unmarshal(body, &response)
+	if err != nil {
+		logger.Err("Failed to parse reponse body of POST request to /drain endpoint: %s", err)
+		return nil, err
+	}
+	return &response, nil
+}
+
+func getDrainStatus(http_client *http.Client, logger *logging.SimpleLogger) (*server.DrainResponse, error) {
+
+	req, err := http.NewRequest("GET", "http://localhost:4141/drain", nil)
+	if err != nil {
+		logger.Err("Failed to create GET request to /drain endpoint: %s", err)
+		return nil, err
+	}
+
+	resp, err := http_client.Do(req)
+	if err != nil {
+		logger.Err("Failed to make GET request to /drain endpoint: %s", err)
+		return nil, err
+	}
+	defer resp.Body.Close()
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		logger.Err("Failed to read reponse body of GET request to /drain endpoint: %s", err)
+		return nil, err
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		logger.Err("Unexpected status code while making GET request to /drain endpoint: ", resp.StatusCode)
+		logger.Info("Response content: %s", string(body))
+		return nil, errors.New("Unexpected status code")
+	}
+
+	var response server.DrainResponse
+	err = json.Unmarshal(body, &response)
+	if err != nil {
+		logger.Err("Failed to parse reponse body of GET request to /drain endpoint: %s", err)
+		return nil, err
+	}
+	return &response, nil
+}

main.go

@@ -35,8 +35,12 @@ func main() {
 	}
 	version := &cmd.VersionCmd{AtlantisVersion: atlantisVersion}
 	testdrive := &cmd.TestdriveCmd{}
+	drainCmd := &cmd.DrainCmd{
+		Logger: logging.NewSimpleLogger("cmd", false, logging.Info),
+	}
 	cmd.RootCmd.AddCommand(server.Init())
 	cmd.RootCmd.AddCommand(version.Init())
 	cmd.RootCmd.AddCommand(testdrive.Init())
+	cmd.RootCmd.AddCommand(drainCmd.Init())
 	cmd.Execute()
 }

server/drain_controller.go

@@ -0,0 +1,84 @@
+package server
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"sync"
+
+	"github.com/runatlantis/atlantis/server/logging"
+)
+
+// DrainController handles all requests relating to Atlantis drainage (to shutdown properly).
+type DrainController struct {
+	Logger                   *logging.SimpleLogger
+	DrainStarted             bool
+	DrainCompleted           bool
+	OngoingOperationsCounter int
+	mutex                    sync.Mutex
+}
+
+type DrainResponse struct {
+	DrainStarted             bool `json:"started"`
+	DrainCompleted           bool `json:"completed"`
+	OngoingOperationsCounter int  `json:"ongoingOperations"`
+}
+
+// Get is the GET /drain route. It renders the current drainage status.
+func (d *DrainController) Get(w http.ResponseWriter, r *http.Request) {
+	d.respondStatus(http.StatusOK, w)
+}
+
+// Post is the POST /drain route. It asks atlantis to finish all ongoing operations and to refuse to start new ones.
+func (d *DrainController) Post(w http.ResponseWriter, r *http.Request) {
+	d.mutex.Lock()
+	defer d.mutex.Unlock()
+	d.DrainStarted = true
+	if d.OngoingOperationsCounter == 0 {
+		d.DrainCompleted = true
+	}
+	d.respondStatus(http.StatusCreated, w)
+}
+
+// Try to add an operation as ongoing. Return true if the operation is allowed to start, false if it should be rejected.
+func (d *DrainController) TryAddNewOngoingOperation() bool {
+	d.mutex.Lock()
+	defer d.mutex.Unlock()
+	if d.DrainStarted {
+		return false
+	} else {
+		d.OngoingOperationsCounter += 1
+		return true
+	}
+}
+
+// Consider on operation as completed.
+func (d *DrainController) RemoveOngoingOperation() {
+	d.mutex.Lock()
+	defer d.mutex.Unlock()
+	d.OngoingOperationsCounter -= 1
+	if d.OngoingOperationsCounter < 0 {
+		d.Logger.Log(logging.Warn, "Drain OngoingOperationsCounter became below 0, this is a bug")
+		d.OngoingOperationsCounter = 0
+	}
+	if d.DrainStarted && d.OngoingOperationsCounter == 0 {
+		d.DrainCompleted = true
+	}
+}
+
+func (d *DrainController) respondStatus(responseCode int, w http.ResponseWriter) {
+	data, err := json.MarshalIndent(&DrainResponse{
+		DrainStarted:             d.DrainStarted,
+		DrainCompleted:           d.DrainCompleted,
+		OngoingOperationsCounter: d.OngoingOperationsCounter,
+	}, "", "  ")
+	if err != nil {
+		w.WriteHeader(http.StatusInternalServerError)
+		fmt.Fprintf(w, "Error creating status json response: %s", err)
+		return
+	}
+	d.Logger.Log(logging.Info, "Drain status: %s", string(data))
+	w.WriteHeader(responseCode)
+	w.Header().Set("Content-Type", "application/json")
+	w.Write(data) // nolint: errcheck
+}

server/drain_controller_test.go

@@ -0,0 +1,390 @@
+package server_test
+
+import (
+	"bytes"
+	"encoding/json"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/runatlantis/atlantis/server"
+	"github.com/runatlantis/atlantis/server/logging"
+	myTests "github.com/runatlantis/atlantis/testing"
+)
+
+func TestDrainController_Get(t *testing.T) {
+	type fields struct {
+		DrainStarted             bool
+		DrainCompleted           bool
+		OngoingOperationsCounter int
+	}
+	type wants struct {
+		Status                   int
+		DrainStarted             bool
+		DrainCompleted           bool
+		OngoingOperationsCounter int
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		wants  wants
+	}{
+		{
+			name: "simple",
+			fields: fields{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 0,
+			},
+			wants: wants{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 0,
+				Status:                   http.StatusOK,
+			},
+		},
+		{
+			name: "on ongoing",
+			fields: fields{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+			},
+			wants: wants{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+				Status:                   http.StatusOK,
+			},
+		},
+		{
+			name: "started",
+			fields: fields{
+				DrainStarted:             true,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 0,
+			},
+			wants: wants{
+				DrainStarted:             true,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 0,
+				Status:                   http.StatusOK,
+			},
+		},
+		{
+			name: "started and completed",
+			fields: fields{
+				DrainStarted:             true,
+				DrainCompleted:           true,
+				OngoingOperationsCounter: 0,
+			},
+			wants: wants{
+				DrainStarted:             true,
+				DrainCompleted:           true,
+				OngoingOperationsCounter: 0,
+				Status:                   http.StatusOK,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			logger := logging.NewNoopLogger()
+			r, _ := http.NewRequest("GET", "/drain", bytes.NewBuffer(nil))
+			w := httptest.NewRecorder()
+			d := &server.DrainController{
+				Logger:                   logger,
+				DrainStarted:             tt.fields.DrainStarted,
+				DrainCompleted:           tt.fields.DrainCompleted,
+				OngoingOperationsCounter: tt.fields.OngoingOperationsCounter,
+			}
+			d.Get(w, r)
+
+			var result server.DrainReponse
+			t.Helper()
+			body, err := ioutil.ReadAll(w.Result().Body)
+			myTests.Ok(t, err)
+			myTests.Assert(t, tt.wants.Status == w.Result().StatusCode, "exp %d got %d, body: %s", tt.wants.Status, w.Result().StatusCode, string(body))
+			err = json.Unmarshal(body, &result)
+			myTests.Ok(t, err)
+			myTests.Assert(t, tt.wants.DrainStarted == result.DrainStarted, "exp %s got %s in DrainStarted of %s", tt.wants.DrainStarted, result.DrainStarted, string(body))
+			myTests.Assert(t, tt.wants.DrainCompleted == result.DrainCompleted, "exp %s got %s in DrainCompleted of %s", tt.wants.DrainCompleted, result.DrainCompleted, string(body))
+			myTests.Assert(t, tt.wants.OngoingOperationsCounter == result.OngoingOperationsCounter, "exp %s got %s in OngoingOperationsCounter of %s", tt.wants.OngoingOperationsCounter, result.OngoingOperationsCounter, string(body))
+		})
+	}
+}
+
+func TestDrainController_Post(t *testing.T) {
+	type fields struct {
+		DrainStarted             bool
+		DrainCompleted           bool
+		OngoingOperationsCounter int
+	}
+	type wants struct {
+		Status                   int
+		DrainStarted             bool
+		DrainCompleted           bool
+		OngoingOperationsCounter int
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		wants  wants
+	}{
+		{
+			name: "simple",
+			fields: fields{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 0,
+			},
+			wants: wants{
+				DrainStarted:             true,
+				DrainCompleted:           true,
+				OngoingOperationsCounter: 0,
+				Status:                   http.StatusCreated,
+			},
+		},
+		{
+			name: "on ongoing",
+			fields: fields{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+			},
+			wants: wants{
+				DrainStarted:             true,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+				Status:                   http.StatusCreated,
+			},
+		},
+		{
+			name: "already started",
+			fields: fields{
+				DrainStarted:             true,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+			},
+			wants: wants{
+				DrainStarted:             true,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+				Status:                   http.StatusCreated,
+			},
+		},
+		{
+			name: "already started and completed",
+			fields: fields{
+				DrainStarted:             true,
+				DrainCompleted:           true,
+				OngoingOperationsCounter: 0,
+			},
+			wants: wants{
+				DrainStarted:             true,
+				DrainCompleted:           true,
+				OngoingOperationsCounter: 0,
+				Status:                   http.StatusCreated,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			logger := logging.NewNoopLogger()
+			r, _ := http.NewRequest("GET", "/drain", bytes.NewBuffer(nil))
+			w := httptest.NewRecorder()
+			d := &server.DrainController{
+				Logger:                   logger,
+				DrainStarted:             tt.fields.DrainStarted,
+				DrainCompleted:           tt.fields.DrainCompleted,
+				OngoingOperationsCounter: tt.fields.OngoingOperationsCounter,
+			}
+			d.Post(w, r)
+
+			var result server.DrainReponse
+			t.Helper()
+			body, err := ioutil.ReadAll(w.Result().Body)
+			myTests.Ok(t, err)
+			myTests.Assert(t, tt.wants.Status == w.Result().StatusCode, "exp %d got %d, body: %s", tt.wants.Status, w.Result().StatusCode, string(body))
+			err = json.Unmarshal(body, &result)
+			myTests.Ok(t, err)
+			myTests.Assert(t, tt.wants.DrainStarted == result.DrainStarted, "exp %s got %s in DrainStarted of %s", tt.wants.DrainStarted, result.DrainStarted, string(body))
+			myTests.Assert(t, tt.wants.DrainCompleted == result.DrainCompleted, "exp %s got %s in DrainCompleted of %s", tt.wants.DrainCompleted, result.DrainCompleted, string(body))
+			myTests.Assert(t, tt.wants.OngoingOperationsCounter == result.OngoingOperationsCounter, "exp %s got %s in OngoingOperationsCounter of %s", tt.wants.OngoingOperationsCounter, result.OngoingOperationsCounter, string(body))
+		})
+	}
+}
+
+func TestDrainController_TryAddNewOngoingOperation(t *testing.T) {
+	type fields struct {
+		DrainStarted             bool
+		DrainCompleted           bool
+		OngoingOperationsCounter int
+	}
+	type wants struct {
+		DrainStarted             bool
+		DrainCompleted           bool
+		OngoingOperationsCounter int
+		Result                   bool
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		wants  wants
+	}{
+		{
+			name: "simple",
+			fields: fields{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 0,
+			},
+			wants: wants{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+				Result:                   true,
+			},
+		},
+		{
+			name: "already started",
+			fields: fields{
+				DrainStarted:             true,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+			},
+			wants: wants{
+				DrainStarted:             true,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+				Result:                   false,
+			},
+		},
+		{
+			name: "already completed",
+			fields: fields{
+				DrainStarted:             true,
+				DrainCompleted:           true,
+				OngoingOperationsCounter: 1,
+			},
+			wants: wants{
+				DrainStarted:             true,
+				DrainCompleted:           true,
+				OngoingOperationsCounter: 1,
+				Result:                   false,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			logger := logging.NewNoopLogger()
+			d := &server.DrainController{
+				Logger:                   logger,
+				DrainStarted:             tt.fields.DrainStarted,
+				DrainCompleted:           tt.fields.DrainCompleted,
+				OngoingOperationsCounter: tt.fields.OngoingOperationsCounter,
+			}
+
+			result := d.TryAddNewOngoingOperation()
+
+			t.Helper()
+			myTests.Assert(t, tt.wants.Result == result, "exp %d got %d", tt.wants.Result, result)
+			myTests.Assert(t, tt.wants.DrainStarted == d.DrainStarted, "exp %s got %s in DrainStarted", tt.wants.DrainStarted, d.DrainStarted)
+			myTests.Assert(t, tt.wants.DrainCompleted == d.DrainCompleted, "exp %s got %s in DrainCompleted", tt.wants.DrainCompleted, d.DrainCompleted)
+			myTests.Assert(t, tt.wants.OngoingOperationsCounter == d.OngoingOperationsCounter, "exp %s got %s in OngoingOperationsCounter", tt.wants.OngoingOperationsCounter, d.OngoingOperationsCounter)
+		})
+	}
+}
+
+func TestDrainController_RemoveOngoingOperation(t *testing.T) {
+	type fields struct {
+		DrainStarted             bool
+		DrainCompleted           bool
+		OngoingOperationsCounter int
+	}
+	type wants struct {
+		DrainStarted             bool
+		DrainCompleted           bool
+		OngoingOperationsCounter int
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		wants  wants
+	}{
+		{
+			name: "simple",
+			fields: fields{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+			},
+			wants: wants{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 0,
+			},
+		},
+		{
+			name: "already started",
+			fields: fields{
+				DrainStarted:             true,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 1,
+			},
+			wants: wants{
+				DrainStarted:             true,
+				DrainCompleted:           true,
+				OngoingOperationsCounter: 0,
+			},
+		},
+		{
+			name: "going negative - not started",
+			fields: fields{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 0,
+			},
+			wants: wants{
+				DrainStarted:             false,
+				DrainCompleted:           false,
+				OngoingOperationsCounter: 0,
+			},
+		},
+		{
+			name: "going negative - started",
+			fields: fields{
+				DrainStarted:             true,
+				DrainCompleted:           true,
+				OngoingOperationsCounter: 0,
+			},
+			wants: wants{
+				DrainStarted:             true,
+				DrainCompleted:           true,
+				OngoingOperationsCounter: 0,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			logger := logging.NewNoopLogger()
+			d := &server.DrainController{
+				Logger:                   logger,
+				DrainStarted:             tt.fields.DrainStarted,
+				DrainCompleted:           tt.fields.DrainCompleted,
+				OngoingOperationsCounter: tt.fields.OngoingOperationsCounter,
+			}
+
+			d.RemoveOngoingOperation()
+
+			t.Helper()
+			myTests.Assert(t, tt.wants.DrainStarted == d.DrainStarted, "exp %s got %s in DrainStarted", tt.wants.DrainStarted, d.DrainStarted)
+			myTests.Assert(t, tt.wants.DrainCompleted == d.DrainCompleted, "exp %s got %s in DrainCompleted", tt.wants.DrainCompleted, d.DrainCompleted)
+			myTests.Assert(t, tt.wants.OngoingOperationsCounter == d.OngoingOperationsCounter, "exp %s got %s in OngoingOperationsCounter", tt.wants.OngoingOperationsCounter, d.OngoingOperationsCounter)
+		})
+	}
+}

server/events_controller.go

@@ -83,6 +83,7 @@ type EventsController struct {
 	// Azure DevOps Team Project. If empty, no request validation is done.
 	AzureDevopsWebhookBasicPassword []byte
 	AzureDevopsRequestValidator     AzureDevopsRequestValidator
+	DrainController                 *DrainController
 }
 
 // Post handles POST webhook requests.
@@ -319,6 +320,17 @@ func (e *EventsController) HandleGithubPullRequestEvent(w http.ResponseWriter, p
 }
 
 func (e *EventsController) handlePullRequestEvent(w http.ResponseWriter, baseRepo models.Repo, headRepo models.Repo, pull models.PullRequest, user models.User, eventType models.PullRequestEventType) {
+	if canProceed := e.DrainController.TryAddNewOngoingOperation(); !canProceed {
+		if commentErr := e.VCSClient.CreateComment(baseRepo, pull.Num, "Atlantis server is shutting down, please try again later."); commentErr != nil {
+			e.Logger.Log(logging.Error, "unable to comment: %s", commentErr)
+		}
+		return
+	}
+
+	defer func() {
+		e.DrainController.RemoveOngoingOperation()
+	}()
+
 	if !e.RepoWhitelistChecker.IsWhitelisted(baseRepo.FullName, baseRepo.VCSHost.Hostname) {
 		// If the repo isn't whitelisted and we receive an opened pull request
 		// event we comment back on the pull request that the repo isn't
@@ -403,6 +415,17 @@ func (e *EventsController) HandleGitlabCommentEvent(w http.ResponseWriter, event
 }
 
 func (e *EventsController) handleCommentEvent(w http.ResponseWriter, baseRepo models.Repo, maybeHeadRepo *models.Repo, maybePull *models.PullRequest, user models.User, pullNum int, comment string, vcsHost models.VCSHostType) {
+	if canProceed := e.DrainController.TryAddNewOngoingOperation(); !canProceed {
+		if commentErr := e.VCSClient.CreateComment(baseRepo, pullNum, "Atlantis server is shutting down, please try again later."); commentErr != nil {
+			e.Logger.Log(logging.Error, "unable to comment: %s", commentErr)
+		}
+		return
+	}
+
+	defer func() {
+		e.DrainController.RemoveOngoingOperation()
+	}()
+
 	parseResult := e.CommentParser.Parse(comment, vcsHost)
 	if parseResult.Ignore {
 		truncated := comment

server/server.go

@@ -75,6 +75,7 @@ type Server struct {
 	Locker             locking.Locker
 	EventsController   *EventsController
 	LocksController    *LocksController
+	DrainController    *DrainController
 	IndexTemplate      TemplateWriter
 	LockDetailTemplate TemplateWriter
 	SSLCertFile        string
@@ -378,6 +379,9 @@ func NewServer(userConfig UserConfig, config Config) (*Server, error) {
 		WorkingDirLocker:   workingDirLocker,
 		DB:                 boltdb,
 	}
+	drainController := &DrainController{
+		Logger: logger,
+	}
 	eventsController := &EventsController{
 		CommandRunner:                   commandRunner,
 		PullCleaner:                     pullClosedExecutor,
@@ -396,6 +400,7 @@ func NewServer(userConfig UserConfig, config Config) (*Server, error) {
 		AzureDevopsWebhookBasicUser:     []byte(userConfig.AzureDevopsWebhookUser),
 		AzureDevopsWebhookBasicPassword: []byte(userConfig.AzureDevopsWebhookPassword),
 		AzureDevopsRequestValidator:     &DefaultAzureDevopsRequestValidator{},
+		DrainController:                 drainController,
 	}
 	return &Server{
 		AtlantisVersion:    config.AtlantisVersion,
@@ -407,6 +412,7 @@ func NewServer(userConfig UserConfig, config Config) (*Server, error) {
 		Locker:             lockingClient,
 		EventsController:   eventsController,
 		LocksController:    locksController,
+		DrainController:    drainController,
 		IndexTemplate:      indexTemplate,
 		LockDetailTemplate: lockTemplate,
 		SSLKeyFile:         userConfig.SSLKeyFile,
@@ -420,6 +426,8 @@ func (s *Server) Start() error {
 		return r.URL.Path == "/" || r.URL.Path == "/index.html"
 	})
 	s.Router.HandleFunc("/healthz", s.Healthz).Methods("GET")
+	s.Router.HandleFunc("/drain", s.DrainController.Get).Methods("GET")
+	s.Router.HandleFunc("/drain", s.DrainController.Post).Methods("POST")
 	s.Router.PathPrefix("/static/").Handler(http.FileServer(&assetfs.AssetFS{Asset: static.Asset, AssetDir: static.AssetDir, AssetInfo: static.AssetInfo}))
 	s.Router.HandleFunc("/events", s.EventsController.Post).Methods("POST")
 	s.Router.HandleFunc("/locks", s.LocksController.DeleteLock).Methods("DELETE").Queries("id", "{id:.*}")