X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/6fe6390690471cee8ba23984e3560fc4ced8b180..153ffc8536bd1f3fdcd1901fbbb7453ef2fcf2e2:/services/keepproxy/keepproxy.go

diff --git a/services/keepproxy/keepproxy.go b/services/keepproxy/keepproxy.go
index 604e93c0fe..3d1b447625 100644
--- a/services/keepproxy/keepproxy.go
+++ b/services/keepproxy/keepproxy.go
@@ -1,3 +1,7 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package main
 
 import (
@@ -20,6 +24,7 @@ import (
 	"git.curoverse.com/arvados.git/sdk/go/arvados"
 	"git.curoverse.com/arvados.git/sdk/go/arvadosclient"
 	"git.curoverse.com/arvados.git/sdk/go/config"
+	"git.curoverse.com/arvados.git/sdk/go/health"
 	"git.curoverse.com/arvados.git/sdk/go/keepclient"
 	"github.com/coreos/go-systemd/daemon"
 	"github.com/ghodss/yaml"
@@ -35,6 +40,7 @@ type Config struct {
 	Timeout         arvados.Duration
 	PIDFile         string
 	Debug           bool
+	ManagementToken string
 }
 
 func DefaultConfig() *Config {
@@ -62,6 +68,7 @@ func main() {
 	flagset.IntVar(&cfg.DefaultReplicas, "default-replicas", cfg.DefaultReplicas, "Default number of replicas to write if not specified by the client. If 0, use site default."+deprecated)
 	flagset.StringVar(&cfg.PIDFile, "pid", cfg.PIDFile, "Path to write pid file."+deprecated)
 	timeoutSeconds := flagset.Int("timeout", int(time.Duration(cfg.Timeout)/time.Second), "Timeout (in seconds) on requests to internal Keep services."+deprecated)
+	flagset.StringVar(&cfg.ManagementToken, "management-token", cfg.ManagementToken, "Authorization token to be included in all health check requests.")
 
 	var cfgPath string
 	const defaultCfgPath = "/etc/arvados/keepproxy/keepproxy.yml"
@@ -104,6 +111,7 @@ func main() {
 	if err != nil {
 		log.Fatalf("Error setting up keep client %s", err.Error())
 	}
+	keepclient.RefreshServiceDiscoveryOnSIGHUP()
 
 	if cfg.PIDFile != "" {
 		f, err := os.Create(cfg.PIDFile)
@@ -133,7 +141,6 @@ func main() {
 	if cfg.DefaultReplicas > 0 {
 		kc.Want_replicas = cfg.DefaultReplicas
 	}
-	go kc.RefreshServices(5*time.Minute, 3*time.Second)
 
 	listener, err = net.Listen("tcp", cfg.Listen)
 	if err != nil {
@@ -156,7 +163,7 @@ func main() {
 	signal.Notify(term, syscall.SIGINT)
 
 	// Start serving requests.
-	router = MakeRESTRouter(!cfg.DisableGet, !cfg.DisablePut, kc, time.Duration(cfg.Timeout))
+	router = MakeRESTRouter(!cfg.DisableGet, !cfg.DisablePut, kc, time.Duration(cfg.Timeout), cfg.ManagementToken)
 	http.Serve(listener, router)
 
 	log.Println("shutting down")
@@ -246,32 +253,29 @@ type proxyHandler struct {
 
 // MakeRESTRouter returns an http.Handler that passes GET and PUT
 // requests to the appropriate handlers.
-func MakeRESTRouter(enable_get bool, enable_put bool, kc *keepclient.KeepClient, timeout time.Duration) http.Handler {
+func MakeRESTRouter(enable_get bool, enable_put bool, kc *keepclient.KeepClient, timeout time.Duration, mgmtToken string) http.Handler {
 	rest := mux.NewRouter()
+
+	transport := *(http.DefaultTransport.(*http.Transport))
+	transport.DialContext = (&net.Dialer{
+		Timeout:   keepclient.DefaultConnectTimeout,
+		KeepAlive: keepclient.DefaultKeepAlive,
+		DualStack: true,
+	}).DialContext
+	transport.TLSClientConfig = arvadosclient.MakeTLSConfig(kc.Arvados.ApiInsecure)
+	transport.TLSHandshakeTimeout = keepclient.DefaultTLSHandshakeTimeout
+
 	h := &proxyHandler{
 		Handler:    rest,
 		KeepClient: kc,
 		timeout:    timeout,
-		transport: &http.Transport{
-			Dial: (&net.Dialer{
-				Timeout:   20 * time.Second,
-				KeepAlive: 10 * time.Second,
-			}).Dial,
-			TLSClientConfig:     arvadosclient.MakeTLSConfig(kc.Arvados.ApiInsecure),
-			TLSHandshakeTimeout: 10 * time.Second,
-		},
+		transport:  &transport,
 		ApiTokenCache: &ApiTokenCache{
 			tokens:     make(map[string]int64),
 			expireTime: 300,
 		},
 	}
 
-	go func(t *http.Transport) {
-		for range time.NewTicker(5 * time.Minute).C {
-			t.CloseIdleConnections()
-		}
-	}(h.transport)
-
 	if enable_get {
 		rest.HandleFunc(`/{locator:[0-9a-f]{32}\+.*}`, h.Get).Methods("GET", "HEAD")
 		rest.HandleFunc(`/{locator:[0-9a-f]{32}}`, h.Get).Methods("GET", "HEAD")
@@ -291,6 +295,11 @@ func MakeRESTRouter(enable_get bool, enable_put bool, kc *keepclient.KeepClient,
 		rest.HandleFunc(`/`, h.Options).Methods("OPTIONS")
 	}
 
+	rest.Handle("/_health/{check}", &health.Handler{
+		Token:  mgmtToken,
+		Prefix: "/_health/",
+	}).Methods("GET")
+
 	rest.NotFoundHandler = InvalidPathHandler{}
 	return h
 }