diff --git a/README.md b/README.md index 85d8c1fd5..f4a3cd8dc 100644 --- a/README.md +++ b/README.md @@ -50,6 +50,7 @@ GLOBAL OPTIONS: --tls_enabled This flag has been deprecated. Specify tls_cert_file and tls_key_file instead. [$BAZEL_REMOTE_TLS_ENABLED] --tls_cert_file value Path to a pem encoded certificate file. [$BAZEL_REMOTE_TLS_CERT_FILE] --tls_key_file value Path to a pem encoded key file. [$BAZEL_REMOTE_TLS_KEY_FILE] + --idle_timeout value The maximum period of having received no request after which the server will shut itself down. Disabled by default. (default: 0s) [$BAZEL_REMOTE_IDLE_TIMEOUT] --help, -h show help ``` diff --git a/config/config.go b/config/config.go index 70bbad43a..ca8d5f2a6 100644 --- a/config/config.go +++ b/config/config.go @@ -5,6 +5,7 @@ import ( "fmt" "io/ioutil" "os" + "time" yaml "gopkg.in/yaml.v2" ) @@ -30,11 +31,12 @@ type Config struct { TLSKeyFile string `yaml:"tls_key_file"` GoogleCloudStorage *GoogleCloudStorageConfig `yaml:"gcs_proxy"` HTTPBackend *HTTPBackendConfig `yaml:"http_proxy"` + IdleTimeout time.Duration `yaml:"idle_timeout"` } // New ... func New(dir string, maxSize int, host string, port int, htpasswdFile string, - tlsCertFile string, tlsKeyFile string) (*Config, error) { + tlsCertFile string, tlsKeyFile string, idleTimeout time.Duration) (*Config, error) { c := Config{ Host: host, Port: port, @@ -45,6 +47,7 @@ func New(dir string, maxSize int, host string, port int, htpasswdFile string, TLSKeyFile: tlsKeyFile, GoogleCloudStorage: nil, HTTPBackend: nil, + IdleTimeout: idleTimeout, } err := validateConfig(&c) diff --git a/main.go b/main.go index 7143eef98..6ad28b2cc 100644 --- a/main.go +++ b/main.go @@ -1,12 +1,15 @@ package main import ( + "context" "fmt" "log" "net/http" "net/url" "os" "strconv" + "sync" + "time" auth "github.com/abbot/go-http-auth" "github.com/buchgr/bazel-remote/cache" @@ -81,6 +84,12 @@ func main() { Usage: "Path to a pem encoded key file.", EnvVar: "BAZEL_REMOTE_TLS_KEY_FILE", }, + cli.DurationFlag{ + Name: "idle_timeout", + Value: 0, + Usage: "The maximum period of having received no request after which the server will shut itself down. Disabled by default.", + EnvVar: "BAZEL_REMOTE_IDLE_TIMEOUT", + }, } app.Action = func(ctx *cli.Context) error { @@ -96,7 +105,8 @@ func main() { ctx.Int("port"), ctx.String("htpasswd_file"), ctx.String("tls_cert_file"), - ctx.String("tls_key_file")) + ctx.String("tls_key_file"), + ctx.Duration("idle_timeout")) } if err != nil { @@ -130,29 +140,65 @@ func main() { proxyCache = diskCache } + mux := http.NewServeMux() + httpServer := &http.Server{ + Addr: c.Host + ":" + strconv.Itoa(c.Port), + Handler: mux, + } h := server.NewHTTPCache(proxyCache, accessLogger, errorLogger) + mux.HandleFunc("/status", h.StatusPageHandler) - http.HandleFunc("/status", h.StatusPageHandler) - http.HandleFunc("/", maybeAuth(h.CacheHandler, c.HtpasswdFile, c.Host)) + cacheHandler := h.CacheHandler + if c.HtpasswdFile != "" { + cacheHandler = wrapAuthHandler(cacheHandler, c.HtpasswdFile, c.Host) + } + if c.IdleTimeout > 0 { + cacheHandler = wrapIdleHandler(cacheHandler, c.IdleTimeout, accessLogger, httpServer) + } + mux.HandleFunc("/", cacheHandler) if len(c.TLSCertFile) > 0 && len(c.TLSKeyFile) > 0 { - return http.ListenAndServeTLS(c.Host+":"+strconv.Itoa(c.Port), c.TLSCertFile, - c.TLSKeyFile, nil) + return httpServer.ListenAndServeTLS(c.TLSCertFile, c.TLSKeyFile) } - return http.ListenAndServe(c.Host+":"+strconv.Itoa(c.Port), nil) + return httpServer.ListenAndServe() } serverErr := app.Run(os.Args) if serverErr != nil { - log.Fatal("ListenAndServe: ", serverErr) + log.Fatal("bazel-remote terminated: ", serverErr) } } -func maybeAuth(fn http.HandlerFunc, htpasswdFile string, host string) http.HandlerFunc { - if htpasswdFile != "" { - secrets := auth.HtpasswdFileProvider(htpasswdFile) - authenticator := auth.NewBasicAuthenticator(host, secrets) - return auth.JustCheck(authenticator, fn) - } - return fn +func wrapIdleHandler(handler http.HandlerFunc, idleTimeout time.Duration, accessLogger cache.Logger, httpServer *http.Server) http.HandlerFunc { + lastRequest := time.Now() + ticker := time.NewTicker(time.Second) + var m sync.Mutex + go func() { + for { + select { + case now := <-ticker.C: + m.Lock() + elapsed := now.Sub(lastRequest) + m.Unlock() + if elapsed > idleTimeout { + ticker.Stop() + accessLogger.Printf("Shutting down server after having been idle for %v", idleTimeout) + httpServer.Shutdown(context.Background()) + } + } + } + }() + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + now := time.Now() + m.Lock() + lastRequest = now + m.Unlock() + handler(w, r) + }) +} + +func wrapAuthHandler(handler http.HandlerFunc, htpasswdFile string, host string) http.HandlerFunc { + secrets := auth.HtpasswdFileProvider(htpasswdFile) + authenticator := auth.NewBasicAuthenticator(host, secrets) + return auth.JustCheck(authenticator, handler) }