feat(ws): release generic WebSocket client and server v0.31.0-preview

app/app.go

@@ -31,6 +31,7 @@ type App struct {
 	closers          []func() error
 	closeOnce        sync.Once
 	closeErr         error
+	components       *lifecycle.Runtime
 	clientIPResolver *httpx.ClientIPResolver
 	responseEncoder  response.Encoder
 	requests         requestDrain
@@ -170,7 +171,8 @@ func (a *App) Use(middleware ...httpx.Middleware) *App {
 }
 
 // OnShutdown registers a cleanup function owned by the application. Cleanup
-// functions run in reverse registration order after all handlers have returned.
+// functions run in reverse registration order after all handlers and registered
+// background stop/drain hooks have finished.
 // Incomplete draining skips these callbacks and returns ErrResourcesRetained.
 // This keeps omni.New as the only public startup entry point while allowing
 // projects to attach databases, loggers and other resources to its lifecycle.
@@ -181,6 +183,33 @@ func (a *App) OnShutdown(cleanup func() error) *App {
 	return a
 }
 
+// OnStop registers a background component's stop-admission hook. Each stop
+// requires an OnDrain with the same name. Hooks run after HTTP handlers finish
+// and before OnShutdown resources are released. Register before serving.
+func (a *App) OnStop(name string, stop func(context.Context) error) *App {
+	a.componentRuntime().OnStop(name, stop)
+	return a
+}
+
+// OnDrain registers proof that a background component, including its callbacks,
+// has finished using shared resources. Incomplete draining retains resources.
+// Generated services may instead register both hooks on their svc Runtime.
+func (a *App) OnDrain(name string, drain func(context.Context) error) *App {
+	a.componentRuntime().OnDrain(name, drain)
+	return a
+}
+
+func (a *App) componentRuntime() *lifecycle.Runtime {
+	if a.components == nil {
+		r, err := lifecycle.NewRuntimeWithConfig(context.Background(), a.config.Lifecycle)
+		if err != nil {
+			panic(err)
+		}
+		a.components = r
+	}
+	return a.components
+}
+
 // Close stops admission, cancels active requests and waits for their handlers
 // to return before releasing resources. It is idempotent. If cancellation grace
 // expires, ErrResourcesRetained is permanent and no cleanup callbacks run.
@@ -192,7 +221,11 @@ func (a *App) Close() error {
 		if timeout <= 0 {
 			timeout = 5 * time.Second
 		}
-		if err := a.requests.cancelAndWait(timeout); err != nil {
+		err := a.requests.cancelAndWait(timeout)
+		if a.components != nil {
+			err = a.components.Finish(err)
+		}
+		if err != nil {
 			a.closeErr = err
 			return
 		}