Added PreShutdownHooks to the apiserver

staging/src/k8s.io/apiserver/pkg/server/BUILD

@@ -75,6 +75,7 @@ go_library(
         "//vendor/k8s.io/apimachinery/pkg/runtime:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/runtime/schema:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/runtime/serializer:go_default_library",
+        "//vendor/k8s.io/apimachinery/pkg/util/errors:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/util/runtime:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/util/sets:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/util/validation:go_default_library",

staging/src/k8s.io/apiserver/pkg/server/config.go

@@ -460,6 +460,7 @@ func (c completedConfig) New(name string, delegationTarget DelegationTarget) (*G
 		openAPIConfig: c.OpenAPIConfig,
 
 		postStartHooks:         map[string]postStartHookEntry{},
+		preShutdownHooks:       map[string]preShutdownHookEntry{},
 		disabledPostStartHooks: c.DisabledPostStartHooks,
 
 		healthzChecks: c.HealthzChecks,
@@ -473,8 +474,12 @@ func (c completedConfig) New(name string, delegationTarget DelegationTarget) (*G
 		s.postStartHooks[k] = v
 	}
 
+	for k, v := range delegationTarget.PreShutdownHooks() {
+		s.preShutdownHooks[k] = v
+	}
+
 	genericApiServerHookName := "generic-apiserver-start-informers"
-	if c.SharedInformerFactory != nil && !s.isHookRegistered(genericApiServerHookName) {
+	if c.SharedInformerFactory != nil && !s.isPostStartHookRegistered(genericApiServerHookName) {
 		err := s.AddPostStartHook(genericApiServerHookName, func(context PostStartHookContext) error {
 			c.SharedInformerFactory.Start(context.StopCh)
 			return nil

staging/src/k8s.io/apiserver/pkg/server/genericapiserver.go

@@ -134,6 +134,10 @@ type GenericAPIServer struct {
 	postStartHooksCalled   bool
 	disabledPostStartHooks sets.String
 
+	preShutdownHookLock    sync.Mutex
+	preShutdownHooks       map[string]preShutdownHookEntry
+	preShutdownHooksCalled bool
+
 	// healthz checks
 	healthzLock    sync.Mutex
 	healthzChecks  []healthz.HealthzChecker
@@ -163,6 +167,9 @@ type DelegationTarget interface {
 	// PostStartHooks returns the post-start hooks that need to be combined
 	PostStartHooks() map[string]postStartHookEntry
 
+	// PreShutdownHooks returns the pre-stop hooks that need to be combined
+	PreShutdownHooks() map[string]preShutdownHookEntry
+
 	// HealthzChecks returns the healthz checks that need to be combined
 	HealthzChecks() []healthz.HealthzChecker
 
@@ -180,6 +187,9 @@ func (s *GenericAPIServer) UnprotectedHandler() http.Handler {
 func (s *GenericAPIServer) PostStartHooks() map[string]postStartHookEntry {
 	return s.postStartHooks
 }
+func (s *GenericAPIServer) PreShutdownHooks() map[string]preShutdownHookEntry {
+	return s.preShutdownHooks
+}
 func (s *GenericAPIServer) HealthzChecks() []healthz.HealthzChecker {
 	return s.healthzChecks
 }
@@ -205,6 +215,9 @@ func (s emptyDelegate) UnprotectedHandler() http.Handler {
 func (s emptyDelegate) PostStartHooks() map[string]postStartHookEntry {
 	return map[string]postStartHookEntry{}
 }
+func (s emptyDelegate) PreShutdownHooks() map[string]preShutdownHookEntry {
+	return map[string]preShutdownHookEntry{}
+}
 func (s emptyDelegate) HealthzChecks() []healthz.HealthzChecker {
 	return []healthz.HealthzChecker{}
 }
@@ -264,7 +277,7 @@ func (s preparedGenericAPIServer) Run(stopCh <-chan struct{}) error {
 		s.GenericAPIServer.AuditBackend.Shutdown()
 	}
 
-	return nil
+	return s.RunPreShutdownHooks()
 }
 
 // NonBlockingRun spawns the secure http server. An error is

staging/src/k8s.io/apiserver/pkg/server/hooks.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/golang/glog"
 
+	utilerrors "k8s.io/apimachinery/pkg/util/errors"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	"k8s.io/apiserver/pkg/server/healthz"
 	restclient "k8s.io/client-go/rest"
@@ -39,6 +40,9 @@ import (
 // until it becomes easier to use.
 type PostStartHookFunc func(context PostStartHookContext) error
 
+// PreShutdownHookFunc is a function that can be added to the shutdown logic.
+type PreShutdownHookFunc func() error
+
 // PostStartHookContext provides information about this API server to a PostStartHookFunc
 type PostStartHookContext struct {
 	// LoopbackClientConfig is a config for a privileged loopback connection to the API server
@@ -59,6 +63,10 @@ type postStartHookEntry struct {
 	done chan struct{}
 }
 
+type preShutdownHookEntry struct {
+	hook PreShutdownHookFunc
+}
+
 // AddPostStartHook allows you to add a PostStartHook.
 func (s *GenericAPIServer) AddPostStartHook(name string, hook PostStartHookFunc) error {
 	if len(name) == 0 {
@@ -97,6 +105,37 @@ func (s *GenericAPIServer) AddPostStartHookOrDie(name string, hook PostStartHook
 	}
 }
 
+// AddPreShutdownHook allows you to add a PreShutdownHook.
+func (s *GenericAPIServer) AddPreShutdownHook(name string, hook PreShutdownHookFunc) error {
+	if len(name) == 0 {
+		return fmt.Errorf("missing name")
+	}
+	if hook == nil {
+		return nil
+	}
+
+	s.preShutdownHookLock.Lock()
+	defer s.preShutdownHookLock.Unlock()
+
+	if s.preShutdownHooksCalled {
+		return fmt.Errorf("unable to add %q because PreShutdownHooks have already been called", name)
+	}
+	if _, exists := s.preShutdownHooks[name]; exists {
+		return fmt.Errorf("unable to add %q because it is already registered", name)
+	}
+
+	s.preShutdownHooks[name] = preShutdownHookEntry{hook: hook}
+
+	return nil
+}
+
+// AddPreShutdownHookOrDie allows you to add a PostStartHook, but dies on failure
+func (s *GenericAPIServer) AddPreShutdownHookOrDie(name string, hook PreShutdownHookFunc) {
+	if err := s.AddPreShutdownHook(name, hook); err != nil {
+		glog.Fatalf("Error registering PreShutdownHook %q: %v", name, err)
+	}
+}
+
 // RunPostStartHooks runs the PostStartHooks for the server
 func (s *GenericAPIServer) RunPostStartHooks(stopCh <-chan struct{}) {
 	s.postStartHookLock.Lock()
@@ -113,8 +152,24 @@ func (s *GenericAPIServer) RunPostStartHooks(stopCh <-chan struct{}) {
 	}
 }
 
-// isHookRegistered checks whether a given hook is registered
-func (s *GenericAPIServer) isHookRegistered(name string) bool {
+// RunPreShutdownHooks runs the PreShutdownHooks for the server
+func (s *GenericAPIServer) RunPreShutdownHooks() error {
+	var errorList []error
+
+	s.preShutdownHookLock.Lock()
+	defer s.preShutdownHookLock.Unlock()
+	s.preShutdownHooksCalled = true
+
+	for hookName, hookEntry := range s.preShutdownHooks {
+		if err := runPreShutdownHook(hookName, hookEntry); err != nil {
+			errorList = append(errorList, err)
+		}
+	}
+	return utilerrors.NewAggregate(errorList)
+}
+
+// isPostStartHookRegistered checks whether a given PostStartHook is registered
+func (s *GenericAPIServer) isPostStartHookRegistered(name string) bool {
 	s.postStartHookLock.Lock()
 	defer s.postStartHookLock.Unlock()
 	_, exists := s.postStartHooks[name]
@@ -135,6 +190,19 @@ func runPostStartHook(name string, entry postStartHookEntry, context PostStartHo
 	close(entry.done)
 }
 
+func runPreShutdownHook(name string, entry preShutdownHookEntry) error {
+	var err error
+	func() {
+		// don't let the hook *accidentally* panic and kill the server
+		defer utilruntime.HandleCrash()
+		err = entry.hook()
+	}()
+	if err != nil {
+		return fmt.Errorf("PreShutdownHook %q failed: %v", name, err)
+	}
+	return nil
+}
+
 // postStartHookHealthz implements a healthz check for poststarthooks.  It will return a "hookNotFinished"
 // error until the poststarthook is finished.
 type postStartHookHealthz struct {