Skip to content

Commit

Permalink
Stop proxy scheduler on system exit (#4293)
Browse files Browse the repository at this point in the history
  • Loading branch information
milosgajdos committed May 13, 2024
2 parents c49220d + 062309c commit 56a020f
Show file tree
Hide file tree
Showing 5 changed files with 40 additions and 8 deletions.
8 changes: 8 additions & 0 deletions registry/handlers/app.go
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,14 @@ func (app *App) RegisterHealthChecks(healthRegistries ...*health.Registry) {
}
}

// Shutdown close the underlying registry
func (app *App) Shutdown() error {
if r, ok := app.registry.(proxy.Closer); ok {
return r.Close()
}
return nil
}

// register a handler with the application, by route name. The handler will be
// passed through the application filters and context will be constructed at
// request time.
Expand Down
9 changes: 9 additions & 0 deletions registry/proxy/proxyregistry.go
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,15 @@ func (pr *proxyingRegistry) BlobStatter() distribution.BlobStatter {
return pr.embedded.BlobStatter()
}

type Closer interface {
// Close release all resources used by this object
Close() error
}

func (pr *proxyingRegistry) Close() error {
return pr.scheduler.Stop()
}

// authChallenger encapsulates a request to the upstream to establish credential challenges
type authChallenger interface {
tryEstablishChallenges(context.Context) error
Expand Down
8 changes: 5 additions & 3 deletions registry/proxy/scheduler/scheduler.go
Original file line number Diff line number Diff line change
Expand Up @@ -206,12 +206,13 @@ func (ttles *TTLExpirationScheduler) startTimer(entry *schedulerEntry, ttl time.
}

// Stop stops the scheduler.
func (ttles *TTLExpirationScheduler) Stop() {
func (ttles *TTLExpirationScheduler) Stop() error {
ttles.Lock()
defer ttles.Unlock()

if err := ttles.writeState(); err != nil {
dcontext.GetLogger(ttles.ctx).Errorf("Error writing scheduler state: %s", err)
err := ttles.writeState()
if err != nil {
err = fmt.Errorf("error writing scheduler state: %w", err)
}

for _, entry := range ttles.entries {
Expand All @@ -221,6 +222,7 @@ func (ttles *TTLExpirationScheduler) Stop() {
close(ttles.doneChan)
ttles.saveTimer.Stop()
ttles.stopped = true
return err
}

func (ttles *TTLExpirationScheduler) writeState() error {
Expand Down
12 changes: 10 additions & 2 deletions registry/proxy/scheduler/scheduler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,12 @@ func TestRestoreOld(t *testing.T) {
if err != nil {
t.Fatalf("Error starting ttlExpirationScheduler: %s", err)
}
defer s.Stop()
defer func(s *TTLExpirationScheduler) {
err := s.Stop()
if err != nil {
t.Fatalf("Error stopping ttlExpirationScheduler: %s", err)
}
}(s)

wg.Wait()
mu.Lock()
Expand Down Expand Up @@ -177,7 +182,10 @@ func TestStopRestore(t *testing.T) {

// Start and stop before all operations complete
// state will be written to fs
s.Stop()
err = s.Stop()
if err != nil {
t.Fatalf(err.Error())
}
time.Sleep(10 * time.Millisecond)

// v2 will restore state from fs
Expand Down
11 changes: 8 additions & 3 deletions registry/registry.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"context"
"crypto/tls"
"crypto/x509"
"errors"
"fmt"
"net/http"
"os"
Expand Down Expand Up @@ -312,7 +313,7 @@ func (registry *Registry) ListenAndServe() error {
}

// setup channel to get notified on SIGTERM signal
signal.Notify(registry.quit, syscall.SIGTERM)
signal.Notify(registry.quit, os.Interrupt, syscall.SIGTERM)
serveErr := make(chan error)

// Start serving in goroutine and listen for stop signal in main thread
Expand All @@ -332,9 +333,13 @@ func (registry *Registry) ListenAndServe() error {
}
}

// Shutdown gracefully shuts down the registry's HTTP server.
// Shutdown gracefully shuts down the registry's HTTP server and application object.
func (registry *Registry) Shutdown(ctx context.Context) error {
return registry.server.Shutdown(ctx)
err := registry.server.Shutdown(ctx)
if appErr := registry.app.Shutdown(); appErr != nil {
err = errors.Join(err, appErr)
}
return err
}

func configureDebugServer(config *configuration.Configuration) {
Expand Down

0 comments on commit 56a020f

Please sign in to comment.