servenv: add --onclose_timeout flag

go/vt/servenv/run.go

@@ -69,7 +69,7 @@ func Run(port int) {
 	}
 
 	log.Info("Shutting down gracefully")
-	Close()
+	fireOnCloseHooks(*onCloseTimeout)
 }
 
 // Close runs any registered exit hooks in parallel.

go/vt/servenv/servenv.go

@@ -57,6 +57,7 @@ var (
 	// Flags to alter the behavior of the library.
 	lameduckPeriod = flag.Duration("lameduck-period", 50*time.Millisecond, "keep running at least this long after SIGTERM before stopping")
 	onTermTimeout  = flag.Duration("onterm_timeout", 10*time.Second, "wait no more than this for OnTermSync handlers before stopping")
+	onCloseTimeout = flag.Duration("onclose_timeout", time.Nanosecond, "wait no more than this for OnClose handlers before stopping")
 	_              = flag.Int("mem-profile-rate", 512*1024, "deprecated: use '-pprof=mem' instead")
 	_              = flag.Int("mutex-profile-fraction", 0, "deprecated: use '-pprof=mutex' instead")
 	catchSigpipe   = flag.Bool("catch-sigpipe", false, "catch and ignore SIGPIPE on stdout and stderr if specified")
@@ -165,24 +166,37 @@ func OnTermSync(f func()) {
 
 // fireOnTermSyncHooks returns true iff all the hooks finish before the timeout.
 func fireOnTermSyncHooks(timeout time.Duration) bool {
+	return fireHooksWithTimeout(timeout, "OnTermSync", onTermSyncHooks.Fire)
+}
+
+// fireOnCloseHooks returns true iff all the hooks finish before the timeout.
+func fireOnCloseHooks(timeout time.Duration) bool {
+	return fireHooksWithTimeout(timeout, "OnClose", func() {
+		onCloseHooks.Fire()
+		ListeningURL = url.URL{}
+	})
+}
+
+// fireHooksWithTimeout returns true iff all the hooks finish before the timeout.
+func fireHooksWithTimeout(timeout time.Duration, name string, hookFn func()) bool {
 	defer log.Flush()
-	log.Infof("Firing synchronous OnTermSync hooks and waiting up to %v for them", timeout)
+	log.Infof("Firing %s hooks and waiting up to %v for them", name, timeout)
 
 	timer := time.NewTimer(timeout)
 	defer timer.Stop()
 
 	done := make(chan struct{})
 	go func() {
-		onTermSyncHooks.Fire()
+		hookFn()
 		close(done)
 	}()
 
 	select {
 	case <-done:
-		log.Infof("OnTermSync hooks finished")
+		log.Infof("%s hooks finished", name)
 		return true
 	case <-timer.C:
-		log.Infof("OnTermSync hooks timed out")
+		log.Infof("%s hooks timed out", name)
 		return false
 	}
 }

go/vt/servenv/servenv_test.go

@@ -57,3 +57,17 @@ func TestFireOnTermSyncHooksTimeout(t *testing.T) {
 		t.Errorf("finished = %v, want %v", finished, want)
 	}
 }
+
+func TestFireOnCloseHooksTimeout(t *testing.T) {
+	onCloseHooks = event.Hooks{}
+
+	OnClose(func() {
+		time.Sleep(1 * time.Second)
+	})
+
+	// we deliberatly test the flag to make sure it's not accidently set to a
+	// high value.
+	if finished, want := fireOnCloseHooks(*onCloseTimeout), false; finished != want {
+		t.Errorf("finished = %v, want %v", finished, want)
+	}
+}