feat: add dependency shutdown and configurable service endpoints

cmd/go-omni/apigen/testdata/managed/client.go

@@ -9,16 +9,23 @@ import (
 var Opened, Closed, DownloadClosed atomic.Int32
 var Started = make(chan struct{}, 1)
 var Release = make(chan struct{})
+var Current atomic.Pointer[ServiceContext]
+var FatalDependency = errors.New("test dependency cannot continue")
 
 // This fake client exercises the same lifetime contract as a database/client.
 type Client struct{}
 
 func initClient(ctx context.Context, c *ServiceContext) error {
+	Current.Store(c)
 	Opened.Add(1)
 	c.Client = &Client{}
-	c.OnClose(func() error { Closed.Add(1); return nil })
+	c.OnClose(func() error { Closed.Add(1); c.Logger.Info("test.client.closed"); return nil })
 	if c.Config.FailInit {
 		return errors.New("client initialization failed")
 	}
+	if c.Config.ShutdownInit {
+		c.RequestShutdown(FatalDependency)
+		return nil
+	}
 	return ctx.Err()
 }

cmd/go-omni/apigen/testdata/managed/operations.go

@@ -71,9 +71,10 @@ type countedReader struct{ io.Reader }
 
 func (c *countedReader) Close() error { svc.DownloadClosed.Add(1); return nil }
 func (l *Logic) Slow(context.Context) (*types.Receipt, error) {
+	before := svc.Closed.Load()
 	svc.Started <- struct{}{}
 	<-svc.Release
-	if svc.Closed.Load() != 0 {
+	if svc.Closed.Load() != before {
 		return nil, errors.New("dependency closed before drain")
 	}
 	return &types.Receipt{Bytes: 1}, nil

cmd/go-omni/apigen/testdata/managed/shutdown_test.go

@@ -0,0 +1,231 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+
+	"example.com/managed/internal/config"
+	"example.com/managed/internal/handler"
+	"example.com/managed/internal/svc"
+	"git.yuntuo.hk/k2safe/go-omni/health"
+	"git.yuntuo.hk/k2safe/go-omni/logx"
+)
+
+func TestGeneratedEndpointSettings(t *testing.T) {
+	cfg := config.Default()
+	cfg.Server.Log.Writer = io.Discard
+	cfg.Server.Endpoints.Liveness.Path = "/livez"
+	cfg.Server.Endpoints.Readiness.Path = "/readyz"
+	cfg.Server.Endpoints.OpenAPI.Enabled = false
+	logger, err := logx.New(cfg.Server.Log)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer logger.Close()
+	deps, err := svc.New(context.Background(), logger, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer deps.Close()
+	app, err := handler.NewApp(deps)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer app.Close()
+	for _, path := range []string{"/health/live", "/health/ready", "/openapi.json"} {
+		if out := response(t, app, "GET", path, nil, nil); out.Code != 404 {
+			t.Fatal(path, out.Code)
+		}
+	}
+	for _, path := range []string{"/readyz", "/livez"} {
+		out := response(t, app, "GET", path, nil, nil)
+		var data map[string]any
+		_ = json.Unmarshal(out.Body.Bytes(), &data)
+		if out.Code != 200 || data["code"] != float64(0) || data["msg"] != "success" || data["message"] != nil || data["requestId"] != nil {
+			t.Fatal(path, out.Code, out.Body.String())
+		}
+	}
+	deps.RequestShutdown(errors.New("private shutdown cause"))
+	out := response(t, app, "GET", "/readyz", nil, nil)
+	if out.Code != 503 || strings.Contains(out.Body.String(), "private") {
+		t.Fatal(out.Code, out.Body.String())
+	}
+	if out = response(t, app, "GET", "/livez", nil, nil); out.Code != 200 {
+		t.Fatal("liveness should still work")
+	}
+	// NewApp uses the project's latest configuration, including endpoint switches.
+	deps.Config.Server.Endpoints.Readiness.ResponseMode = "raw"
+	deps.Config.Server.Endpoints.Liveness.Enabled = false
+	deps.Config.Server.Endpoints.OpenAPI.Enabled = true
+	deps.Config.Server.Endpoints.OpenAPI.Path = "/spec.json"
+	app2, err := handler.NewApp(deps)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer app2.Close()
+	out = response(t, app2, "GET", "/readyz", nil, nil)
+	if out.Code != 503 || strings.Contains(out.Body.String(), `"code"`) || !strings.Contains(out.Body.String(), `"status":"not_ready"`) {
+		t.Fatal(out.Body.String())
+	}
+	out = response(t, app2, "GET", "/spec.json", nil, nil)
+	if out.Code != 200 || !strings.Contains(out.Body.String(), `"openapi":`) || strings.Contains(out.Body.String(), `"msg":"操作成功"`) {
+		t.Fatal("OpenAPI was wrapped")
+	}
+	deps.Config.Server.Endpoints.Readiness.Path = "/check"
+	if _, err = handler.NewApp(deps); err == nil || !strings.Contains(err.Error(), "overlaps business") {
+		t.Fatal("collision accepted", err)
+	}
+	deps.Config.Server.Endpoints.Readiness.Enabled = false
+	if a, err := handler.NewApp(deps); err != nil {
+		t.Fatal("disabled endpoint reserves path", err)
+	} else {
+		_ = a.Close()
+	}
+}
+
+func TestGeneratedFailedProbe(t *testing.T) {
+	cfg := config.Default()
+	cfg.Server.Log.Writer = io.Discard
+	logger, _ := logx.New(cfg.Server.Log)
+	defer logger.Close()
+	deps, err := svc.New(context.Background(), logger, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer deps.Close()
+	if err = deps.Health.Add(health.Probe{Name: "database", Check: func(context.Context) error { return errors.New("private connection details") }}); err != nil {
+		t.Fatal(err)
+	}
+	app, err := handler.NewApp(deps)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer app.Close()
+	out := response(t, app, "GET", "/health/ready", nil, nil)
+	if out.Code != 503 || strings.Contains(out.Body.String(), "private") || !strings.Contains(out.Body.String(), `"msg":`) {
+		t.Fatal(out.Code, out.Body.String())
+	}
+}
+
+func TestGeneratedDependencyShutdown(t *testing.T) {
+	for _, transport := range []string{"nethttp", "native"} {
+		t.Run(transport, func(t *testing.T) {
+			listener, err := net.Listen("tcp", "127.0.0.1:0")
+			if err != nil {
+				t.Fatal(err)
+			}
+			address := listener.Addr().String()
+			_ = listener.Close()
+			t.Setenv("GO_OMNI_TRANSPORT", transport)
+			svc.Started = make(chan struct{}, 1)
+			svc.Release = make(chan struct{})
+			before := svc.Closed.Load()
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			done := make(chan error, 1)
+			path := configFile(t, address, false)
+			go func() { done <- run(ctx, path) }()
+			base := "http://" + address
+			waitReady(t, base, done)
+			deps := svc.Current.Load()
+			requestDone := make(chan error, 1)
+			go func() {
+				client := http.Client{Timeout: 3 * time.Second}
+				res, err := client.Get(base + "/slow")
+				if err == nil {
+					_, _ = io.Copy(io.Discard, res.Body)
+					_ = res.Body.Close()
+					if res.StatusCode != 200 {
+						err = io.ErrUnexpectedEOF
+					}
+				}
+				requestDone <- err
+			}()
+			select {
+			case <-svc.Started:
+			case <-time.After(3 * time.Second):
+				t.Fatal("request did not start")
+			}
+			deps.RequestShutdown(svc.FatalDependency)
+			// Inspect the same generated readiness handler during draining, even if
+			// the network listener has already stopped accepting new connections.
+			probe, _ := handler.NewApp(deps)
+			out := httptest.NewRecorder()
+			probe.ServeHTTP(out, httptest.NewRequest("GET", "/health/ready", nil))
+			_ = probe.Close()
+			if out.Code != 503 {
+				t.Fatal("draining remained ready", out.Body.String())
+			}
+			if svc.Closed.Load() != before {
+				t.Fatal("dependency closed before drain")
+			}
+			close(svc.Release)
+			select {
+			case err = <-requestDone:
+				if err != nil {
+					t.Fatal(err)
+				}
+			case <-time.After(3 * time.Second):
+				t.Fatal("request stuck")
+			}
+			select {
+			case err = <-done:
+				if !errors.Is(err, svc.FatalDependency) {
+					t.Fatal("fatal cause lost", err)
+				}
+			case <-time.After(3 * time.Second):
+				t.Fatal("shutdown stuck")
+			}
+			if svc.Closed.Load() != before+1 {
+				t.Fatal("cleanup did not run once")
+			}
+		})
+	}
+}
+
+func TestGeneratedFatalProcessExit(t *testing.T) {
+	if path := os.Getenv("OMNI_TEST_FATAL_CONFIG"); path != "" {
+		os.Args = []string{os.Args[0], "-config", path}
+		main()
+		return
+	}
+	path := filepath.Join(t.TempDir(), "config.json")
+	if err := os.WriteFile(path, []byte(`{"shutdownInit":true,"server":{"address":"127.0.0.1:0"}}`), 0600); err != nil {
+		t.Fatal(err)
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	cmd := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestGeneratedFatalProcessExit$")
+	cmd.Env = append(os.Environ(), "OMNI_TEST_FATAL_CONFIG="+path)
+	output, err := cmd.CombinedOutput()
+	var exit *exec.ExitError
+	if !errors.As(err, &exit) || exit.ExitCode() != 1 {
+		t.Fatal("fatal shutdown must exit 1", err, string(output))
+	}
+	if !strings.Contains(string(output), svc.FatalDependency.Error()) || strings.Count(string(output), "test.client.closed") != 1 {
+		t.Fatal("missing failure or cleanup", string(output))
+	}
+}
+
+func TestGeneratedEndpointConflictCleansUp(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "config.json")
+	if err := os.WriteFile(path, []byte(`{"server":{"address":"127.0.0.1:0","endpoints":{"readiness":{"path":"/check"}}}}`), 0600); err != nil {
+		t.Fatal(err)
+	}
+	before := svc.Closed.Load()
+	err := run(context.Background(), path)
+	if err == nil || !strings.Contains(err.Error(), "overlaps business") || svc.Closed.Load() != before+1 {
+		t.Fatal("collision leaked resources or started serving", err)
+	}
+}