From e64fbad5a36ae00e23dd2d1038432c5cc758d077 Mon Sep 17 00:00:00 2001 From: Benoit BERAUD Date: Tue, 24 Mar 2020 10:25:57 +0100 Subject: [PATCH] Add drain operation + endpoint to clean properly the server before shutdown --- cmd/drain.go | 27 +++ drain/drain.go | 103 +++++++++ main.go | 4 + server/drain_controller.go | 84 +++++++ server/drain_controller_test.go | 390 ++++++++++++++++++++++++++++++++ server/events_controller.go | 23 ++ server/server.go | 8 + 7 files changed, 639 insertions(+) create mode 100644 cmd/drain.go create mode 100644 drain/drain.go create mode 100644 server/drain_controller.go create mode 100644 server/drain_controller_test.go diff --git a/cmd/drain.go b/cmd/drain.go new file mode 100644 index 000000000..8269ddc40 --- /dev/null +++ b/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, + } +} diff --git a/drain/drain.go b/drain/drain.go new file mode 100644 index 000000000..f806446aa --- /dev/null +++ b/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 +} diff --git a/main.go b/main.go index d2aaea00b..800653820 100644 --- a/main.go +++ b/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() } diff --git a/server/drain_controller.go b/server/drain_controller.go new file mode 100644 index 000000000..9e738cc40 --- /dev/null +++ b/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 +} diff --git a/server/drain_controller_test.go b/server/drain_controller_test.go new file mode 100644 index 000000000..bae3f944d --- /dev/null +++ b/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) + }) + } +} diff --git a/server/events_controller.go b/server/events_controller.go index 2b62a2017..fa7a404d0 100644 --- a/server/events_controller.go +++ b/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 diff --git a/server/server.go b/server/server.go index da7fa4618..3b5dafe0a 100644 --- a/server/server.go +++ b/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:.*}")