304 lines
8.3 KiB
Go
304 lines
8.3 KiB
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 正确跳过重启逻辑")
|
|||
|
|
}
|