hotime/graceful_shutdown_test.go

package hotime

import (
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"code.hoteas.com/golang/hotime/log"
)

// getFreePort 获取一个可用的随机端口
func getFreePort() (int, error) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return 0, err
	}
	port := ln.Addr().(*net.TCPAddr).Port
	ln.Close()
	return port, nil
}

// newTestApp 创建一个用于测试的最小 Application，不依赖数据库和配置文件
func newTestApp(port int, handler http.HandlerFunc) *Application {
	app := &Application{}
	app.Log = log.NewLogger(1, "", 10)
	app.Server = &http.Server{
		Addr:    fmt.Sprintf(":%d", port),
		Handler: app,
	}
	// 用自定义 handler 替代完整的路由系统
	app.Handler = handler
	return app
}

// ServeHTTP 已经在 application.go 中定义，
// 但 handler 方法依赖完整初始化，这里我们直接用 http.Handler 接口。
// 因此测试 app 的 ServeHTTP 会先检查 shuttingDown，然后调用 that.handler()。
// 但 that.handler() 需要完整的路由初始化，不适合单元测试。
// 所以我们直接测试核心逻辑：shuttingDown 标志 + http.Server.Shutdown。

// TestServeHTTP_ShuttingDown 测试停机标志置为 true 后新请求返回 503
func TestServeHTTP_ShuttingDown(t *testing.T) {
	port, err := getFreePort()
	if err != nil {
		t.Fatalf("获取空闲端口失败: %v", err)
	}

	var requestHandled atomic.Int32
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestHandled.Add(1)
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("OK"))
	})

	app := &Application{}
	app.Log = log.NewLogger(1, "", 10)
	app.Server = &http.Server{
		Addr: fmt.Sprintf(":%d", port),
	}
	app.Server.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if app.shuttingDown.Load() {
			w.Header().Set("Connection", "close")
			http.Error(w, "Service Unavailable", http.StatusServiceUnavailable)
			return
		}
		handler.ServeHTTP(w, r)
	})

	go app.Server.ListenAndServe()
	time.Sleep(100 * time.Millisecond)

	baseURL := fmt.Sprintf("http://127.0.0.1:%d", port)
	client := &http.Client{Timeout: 5 * time.Second}

	// --- 正常请求应该返回 200 ---
	resp, err := client.Get(baseURL + "/test")
	if err != nil {
		t.Fatalf("正常请求失败: %v", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("期望 200，实际 %d", resp.StatusCode)
	}
	resp.Body.Close()
	t.Logf("[PASS] 正常请求返回 %d", resp.StatusCode)

	if requestHandled.Load() != 1 {
		t.Fatalf("期望 handler 被调用 1 次，实际 %d 次", requestHandled.Load())
	}

	// --- 设置停机标志 ---
	app.shuttingDown.Store(true)
	t.Log("[INFO] shuttingDown 已置为 true")

	// --- 停机后请求应该返回 503 ---
	resp2, err := client.Get(baseURL + "/test2")
	if err != nil {
		t.Fatalf("停机请求失败: %v", err)
	}
	if resp2.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("期望 503，实际 %d", resp2.StatusCode)
	}
	resp2.Body.Close()
	t.Logf("[PASS] 停机后请求返回 %d", resp2.StatusCode)

	// handler 不应被再次调用
	if requestHandled.Load() != 1 {
		t.Fatalf("停机后 handler 不应被调用，但被调用了 %d 次", requestHandled.Load())
	}
	t.Log("[PASS] 停机后 handler 未被调用")

	app.Server.Close()
}

// TestGracefulShutdown_InFlightRequestCompletes 测试优雅停机等待在途请求完成
func TestGracefulShutdown_InFlightRequestCompletes(t *testing.T) {
	port, err := getFreePort()
	if err != nil {
		t.Fatalf("获取空闲端口失败: %v", err)
	}

	slowRequestStarted := make(chan struct{})
	slowRequestDone := make(chan struct{})

	app := &Application{}
	app.Log = log.NewLogger(1, "", 10)
	app.Server = &http.Server{
		Addr: fmt.Sprintf(":%d", port),
	}
	app.Server.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if app.shuttingDown.Load() {
			w.Header().Set("Connection", "close")
			http.Error(w, "Service Unavailable", http.StatusServiceUnavailable)
			return
		}
		if r.URL.Path == "/slow" {
			close(slowRequestStarted)
			time.Sleep(2 * time.Second) // 模拟慢接口
			w.WriteHeader(http.StatusOK)
			w.Write([]byte("SLOW_DONE"))
			close(slowRequestDone)
			return
		}
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("OK"))
	})

	go app.Server.ListenAndServe()
	time.Sleep(100 * time.Millisecond)

	baseURL := fmt.Sprintf("http://127.0.0.1:%d", port)
	client := &http.Client{Timeout: 10 * time.Second}

	// --- 发起慢请求（模拟在途请求）---
	var slowResp *http.Response
	var slowErr error
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		slowResp, slowErr = client.Get(baseURL + "/slow")
	}()

	// 等慢请求开始处理
	<-slowRequestStarted
	t.Log("[INFO] 慢请求已开始处理")

	// --- 在慢请求进行中触发 Shutdown ---
	app.shuttingDown.Store(true)
	t.Log("[INFO] shuttingDown 已置为 true，开始 Shutdown")

	shutdownDone := make(chan struct{})
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		app.Server.Shutdown(ctx)
		close(shutdownDone)
	}()

	// --- 新请求应该返回 503（但因为 Shutdown 不接受新连接了，可能连接被拒绝）---
	// Shutdown 调用后，新的 TCP 连接也会被拒绝，所以这里验证连接级别的拒绝也OK
	resp, err := client.Get(baseURL + "/new")
	if err != nil {
		t.Logf("[PASS] Shutdown 后新请求连接被拒绝（预期行为）: %v", err)
	} else {
		if resp.StatusCode == http.StatusServiceUnavailable {
			t.Logf("[PASS] Shutdown 后新请求返回 503")
		} else {
			t.Logf("[INFO] Shutdown 后新请求返回 %d", resp.StatusCode)
		}
		resp.Body.Close()
	}

	// --- 等待慢请求完成 ---
	wg.Wait()
	if slowErr != nil {
		t.Fatalf("慢请求返回错误（不应中断在途请求）: %v", slowErr)
	}
	body, _ := ioutil.ReadAll(slowResp.Body)
	slowResp.Body.Close()

	if slowResp.StatusCode != http.StatusOK {
		t.Fatalf("慢请求期望 200，实际 %d", slowResp.StatusCode)
	}
	if string(body) != "SLOW_DONE" {
		t.Fatalf("慢请求期望 body='SLOW_DONE'，实际 '%s'", string(body))
	}
	t.Logf("[PASS] 慢请求正常完成，状态=%d，body=%s", slowResp.StatusCode, string(body))

	// --- 等 Shutdown 完成 ---
	<-shutdownDone
	t.Log("[PASS] Server.Shutdown() 在慢请求完成后返回")

	// 确认慢请求确实跑完了
	select {
	case <-slowRequestDone:
		t.Log("[PASS] 慢请求 handler 完整执行完毕")
	default:
		t.Fatal("慢请求 handler 未完整执行")
	}
}

// TestShutdownOnce_MultipleSignals 测试多次触发只执行一次
func TestShutdownOnce_MultipleSignals(t *testing.T) {
	var callCount atomic.Int32
	app := &Application{}
	app.Log = log.NewLogger(1, "", 10)
	app.Server = &http.Server{}

	// 为了避免 os.Exit(0) 导致测试进程退出，
	// 我们直接测试 shutdownOnce + shuttingDown 逻辑
	var once sync.Once
	var shuttingDown atomic.Bool

	mockShutdown := func() {
		once.Do(func() {
			shuttingDown.Store(true)
			callCount.Add(1)
		})
	}

	// 并发触发 10 次
	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			mockShutdown()
		}()
	}
	wg.Wait()

	if callCount.Load() != 1 {
		t.Fatalf("期望 shutdownOnce.Do 只执行 1 次，实际 %d 次", callCount.Load())
	}
	if !shuttingDown.Load() {
		t.Fatal("shuttingDown 应为 true")
	}
	t.Logf("[PASS] 10 次并发触发，shutdownOnce.Do 只执行了 %d 次", callCount.Load())
}

// TestRecoverSkipsDuringShutdown 测试停机中的 panic 不触发重启
func TestRecoverSkipsDuringShutdown(t *testing.T) {
	app := &Application{}
	app.Log = log.NewLogger(1, "", 10)

	var restartCalled atomic.Bool

	// 模拟 Run() 中的 defer recover 逻辑
	testRecover := func() {
		defer func() {
			if err := recover(); err != nil {
				if app.shuttingDown.Load() {
					return // 停机中不重启
				}
				restartCalled.Store(true)
			}
		}()
		panic("test panic")
	}

	// 未停机时，recover 应触发"重启"
	restartCalled.Store(false)
	testRecover()
	if !restartCalled.Load() {
		t.Fatal("未停机时 panic 应触发重启")
	}
	t.Log("[PASS] 未停机时 panic 正确触发重启逻辑")

	// 停机时，recover 不应触发"重启"
	app.shuttingDown.Store(true)
	restartCalled.Store(false)
	testRecover()
	if restartCalled.Load() {
		t.Fatal("停机时 panic 不应触发重启")
	}
	t.Log("[PASS] 停机时 panic 正确跳过重启逻辑")
}