feat: drain readiness before graceful shutdown

cmd/orbit/main.go

@@ -54,6 +54,12 @@ func run() error {
 		return err
 	}
 	cancel()
+	readiness := httpserver.NewReadiness(func(ctx context.Context) error {
+		if err := db.PingContext(ctx); err != nil {
+			return err
+		}
+		return mysql.CheckSchema(ctx, db)
+	})
 	var metrics *observability.Registry
 	var metricsListener net.Listener
 	if cfg.MetricsAddr != "" {
@@ -71,13 +77,7 @@ func run() error {
 		})
 		metrics.RecordDatabase(db.Stats())
 	}
-	ready := func(ctx context.Context) error {
-		if err := db.PingContext(ctx); err != nil {
-			return err
-		}
-		return mysql.CheckSchema(ctx, db)
-	}
-	server := httpserver.New(cfg.HTTPAddr, ready, service.New(mysql.NewTaskStore(db)))
+	server := httpserver.New(cfg.HTTPAddr, readiness.Check, service.New(mysql.NewTaskStore(db)))
 	httpserver.RegisterApplicationResources(server, resources.New(mysql.NewTaskStore(db)))
 	node, err := core.NewID()
 	if err != nil {
@@ -257,6 +257,7 @@ func run() error {
 		}
 		return err
 	case <-root.Done():
+		readiness.BeginDrain()
 		shutdown, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 		defer cancel()
 		if err := server.Shutdown(shutdown); err != nil {