go-omni/app/lifecycle_test.go

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)
		}
	}
}