package app import ( "context" "errors" "io" "net" "net/http" "sync" "testing" "time" "git.yuntuo.hk/k2safe/go-omni/config" "git.yuntuo.hk/k2safe/go-omni/httpx" "git.yuntuo.hk/k2safe/go-omni/transport/nethttp" ) type closingListener struct { net.Listener closed chan struct{} once sync.Once } func (l *closingListener) Close() error { err := l.Listener.Close() l.once.Do(func() { close(l.closed) }) return err } func startTestServer(t *testing.T, application *App) (*closingListener, context.CancelFunc, <-chan error) { t.Helper() listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatal(err) } wrapped := &closingListener{Listener: listener, closed: make(chan struct{})} ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) t.Cleanup(func() { _ = listener.Close() }) done := make(chan error, 1) go func() { done <- application.serveContext(ctx, application.NewServer(""), wrapped) }() return wrapped, cancel, done } func await(t *testing.T, event <-chan struct{}) { t.Helper() select { case <-event: case <-time.After(5 * time.Second): t.Fatal("timed out waiting for server event") } } func TestShutdownDrainsRequestBeforeCleanup(t *testing.T) { application := NewWithRouter(nethttp.New(), config.WithShutdownTimeout(3*time.Second)) started, release, cleaned := make(chan struct{}), make(chan struct{}), make(chan struct{}) var releaseOnce sync.Once unblock := func() { releaseOnce.Do(func() { close(release) }) } t.Cleanup(unblock) application.GET("/slow", func(ctx *httpx.Context) error { close(started) <-release return ctx.OK("finished") }) application.OnShutdown(func() error { close(cleaned); return nil }) listener, cancel, done := startTestServer(t, application) requestDone := make(chan error, 1) go func() { client := &http.Client{Timeout: 5 * time.Second} response, err := client.Get("http://" + listener.Addr().String() + "/slow") if err == nil { defer response.Body.Close() _, err = io.ReadAll(response.Body) if response.StatusCode != http.StatusOK { err = errors.New("request was not completed successfully") } } requestDone <- err }() await(t, started) cancel() await(t, listener.closed) select { case <-cleaned: t.Fatal("resources closed while a request was still running") default: } unblock() if err := <-requestDone; err != nil { t.Fatal(err) } select { case err := <-done: if err != nil { t.Fatal(err) } case <-time.After(5 * time.Second): t.Fatal("server did not finish shutdown") } await(t, cleaned) } func TestShutdownTimeoutClosesActiveConnection(t *testing.T) { application := NewWithRouter(nethttp.New(), config.WithShutdownTimeout(25*time.Millisecond)) started, canceled := make(chan struct{}), make(chan struct{}) application.GET("/slow", func(ctx *httpx.Context) error { close(started) <-ctx.Context().Done() close(canceled) return nil }) listener, cancel, done := startTestServer(t, application) requestDone := make(chan struct{}) go func() { defer close(requestDone) client := &http.Client{Timeout: 5 * time.Second} response, err := client.Get("http://" + listener.Addr().String() + "/slow") if err == nil { response.Body.Close() } }() await(t, started) cancel() select { case err := <-done: if !errors.Is(err, context.DeadlineExceeded) { t.Fatalf("expected shutdown deadline error, got %v", err) } case <-time.After(5 * time.Second): t.Fatal("shutdown ignored timeout") } await(t, canceled) await(t, requestDone) } func TestRunContextStartupFailureAndCanceledContextCloseResources(t *testing.T) { listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatal(err) } defer listener.Close() for _, canceled := range []bool{false, true} { application := NewWithRouter(nethttp.New()) closeCount := 0 application.OnShutdown(func() error { closeCount++; return nil }) ctx, cancel := context.WithCancel(context.Background()) if canceled { cancel() } err := application.RunContext(ctx, listener.Addr().String()) cancel() if !canceled && err == nil { t.Fatal("expected occupied address error") } if canceled && err != nil { t.Fatalf("canceled start: %v", err) } _ = application.Close() if closeCount != 1 { t.Fatalf("cleanup ran %d times", closeCount) } } }