mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
Merge pull request #4214 from bestony/agent/devbox-coding/25c66071-1783957460
feat: add opt-in Server-Timing for Admin UI APIs
This commit is contained in:
@@ -601,6 +601,7 @@ type ServerConfig struct {
|
||||
Host string `mapstructure:"host"`
|
||||
Port int `mapstructure:"port"`
|
||||
Mode string `mapstructure:"mode"` // debug/release
|
||||
EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header
|
||||
FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL,用于生成邮件中的外部链接
|
||||
ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时(秒)
|
||||
IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时(秒)
|
||||
@@ -1459,6 +1460,9 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
|
||||
// 环境变量支持
|
||||
viper.AutomaticEnv()
|
||||
viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
|
||||
if err := viper.BindEnv("server.enable_server_timing", "ENABLE_SERVER_TIMING"); err != nil {
|
||||
return nil, fmt.Errorf("bind ENABLE_SERVER_TIMING: %w", err)
|
||||
}
|
||||
|
||||
// 默认值
|
||||
setDefaults()
|
||||
@@ -1614,6 +1618,7 @@ func setDefaults() {
|
||||
viper.SetDefault("server.host", "0.0.0.0")
|
||||
viper.SetDefault("server.port", 8080)
|
||||
viper.SetDefault("server.mode", "release")
|
||||
viper.SetDefault("server.enable_server_timing", false)
|
||||
viper.SetDefault("server.frontend_url", "")
|
||||
viper.SetDefault("server.read_header_timeout", 30) // 30秒读取请求头
|
||||
viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时
|
||||
|
||||
@@ -17,6 +17,23 @@ func resetViperWithJWTSecret(t *testing.T) {
|
||||
t.Setenv("JWT_SECRET", strings.Repeat("x", 32))
|
||||
}
|
||||
|
||||
func TestLoadServerTimingConfig(t *testing.T) {
|
||||
t.Run("disabled by default", func(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.False(t, cfg.Server.EnableServerTiming)
|
||||
})
|
||||
|
||||
t.Run("enabled by exact environment variable", func(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("ENABLE_SERVER_TIMING", "true")
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.True(t, cfg.Server.EnableServerTiming)
|
||||
})
|
||||
}
|
||||
|
||||
func TestLoadForBootstrapAllowsMissingJWTSecret(t *testing.T) {
|
||||
viper.Reset()
|
||||
t.Setenv("JWT_SECRET", "")
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
@@ -323,7 +324,7 @@ func (h *OpsHandler) QPSWSHandler(c *gin.Context) {
|
||||
// If realtime monitoring is disabled, prefer a successful WS upgrade followed by a clean close
|
||||
// with a deterministic close code. This prevents clients from spinning on 404/1006 reconnect loops.
|
||||
if !h.opsService.IsRealtimeMonitoringEnabled(c.Request.Context()) {
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, servermiddleware.ServerTimingResponseHeader(c))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "ops realtime monitoring is disabled"})
|
||||
return
|
||||
@@ -358,7 +359,7 @@ func (h *OpsHandler) QPSWSHandler(c *gin.Context) {
|
||||
defer releaseOpsWSIPSlot(clientIP)
|
||||
}
|
||||
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, servermiddleware.ServerTimingResponseHeader(c))
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("handler.admin.ops_ws", "[OpsWS] upgrade failed: %v", err)
|
||||
return
|
||||
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
)
|
||||
|
||||
// ForbiddenError 表示上游返回 403 Forbidden
|
||||
@@ -279,7 +280,6 @@ func NewClient(proxyURL string) (*Client, error) {
|
||||
}
|
||||
client.Transport = transport
|
||||
}
|
||||
|
||||
return &Client{
|
||||
httpClient: client,
|
||||
}, nil
|
||||
@@ -341,7 +341,7 @@ func (c *Client) ExchangeCode(ctx context.Context, code, codeVerifier string) (*
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
resp, err := servertiming.Do(c.httpClient, req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("token 交换请求失败: %w", err)
|
||||
}
|
||||
@@ -383,7 +383,7 @@ func (c *Client) RefreshToken(ctx context.Context, refreshToken string) (*TokenR
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
resp, err := servertiming.Do(c.httpClient, req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("token 刷新请求失败: %w", err)
|
||||
}
|
||||
@@ -414,7 +414,7 @@ func (c *Client) GetUserInfo(ctx context.Context, accessToken string) (*UserInfo
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
resp, err := servertiming.Do(c.httpClient, req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("用户信息请求失败: %w", err)
|
||||
}
|
||||
@@ -465,7 +465,7 @@ func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadC
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
resp, err := servertiming.Do(c.httpClient, req)
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("loadCodeAssist 请求失败: %w", err)
|
||||
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
|
||||
@@ -544,7 +544,7 @@ func (c *Client) OnboardUser(ctx context.Context, accessToken, tierID string) (s
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
resp, err := servertiming.Do(c.httpClient, req)
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("onboardUser 请求失败: %w", err)
|
||||
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
|
||||
@@ -683,7 +683,7 @@ func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectI
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
|
||||
|
||||
resp, err := fetchClient.Do(req)
|
||||
resp, err := servertiming.Do(fetchClient, req)
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("fetchAvailableModels 请求失败: %w", err)
|
||||
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
|
||||
@@ -842,7 +842,7 @@ func (c *Client) SetUserSettings(ctx context.Context, accessToken string) (*SetU
|
||||
req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
|
||||
req.Host = "daily-cloudcode-pa.googleapis.com"
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
resp, err := servertiming.Do(c.httpClient, req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("setUserSettings 请求失败: %w", err)
|
||||
}
|
||||
@@ -885,7 +885,7 @@ func (c *Client) FetchUserInfo(ctx context.Context, accessToken, projectID strin
|
||||
req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
|
||||
req.Host = "daily-cloudcode-pa.googleapis.com"
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
resp, err := servertiming.Do(c.httpClient, req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("fetchUserInfo 请求失败: %w", err)
|
||||
}
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
|
||||
)
|
||||
|
||||
@@ -92,6 +93,7 @@ func buildClient(opts Options) (*http.Client, error) {
|
||||
if opts.ValidateResolvedIP && !opts.AllowPrivateHosts {
|
||||
rt = newValidatedTransport(transport)
|
||||
}
|
||||
rt = servertiming.WrapRoundTripper(rt)
|
||||
return &http.Client{
|
||||
Transport: rt,
|
||||
Timeout: opts.Timeout,
|
||||
|
||||
@@ -0,0 +1,348 @@
|
||||
package servertiming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
HeaderName = "Server-Timing"
|
||||
AdminUIHeader = "X-Admin-UI-Request"
|
||||
MetricDatabase = "db"
|
||||
MetricRedis = "redis"
|
||||
dependencyPrefix = "dep_"
|
||||
|
||||
maxMetricNameLength = 48
|
||||
maxIntervals = 2048
|
||||
maxHeaderLength = 4096
|
||||
)
|
||||
|
||||
type contextKey struct{}
|
||||
|
||||
type interval struct {
|
||||
start time.Time
|
||||
end time.Time
|
||||
}
|
||||
|
||||
type metric struct {
|
||||
count int64
|
||||
intervals []interval
|
||||
}
|
||||
|
||||
// Collector stores request-scoped timing samples. It is safe for concurrent use.
|
||||
type Collector struct {
|
||||
startedAt time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
metrics map[string]*metric
|
||||
cacheStatus string
|
||||
}
|
||||
|
||||
// New creates a collector whose total duration starts at startedAt.
|
||||
func New(startedAt time.Time) *Collector {
|
||||
if startedAt.IsZero() {
|
||||
startedAt = time.Now()
|
||||
}
|
||||
return &Collector{
|
||||
startedAt: startedAt,
|
||||
metrics: make(map[string]*metric),
|
||||
}
|
||||
}
|
||||
|
||||
// WithCollector attaches a collector to a context.
|
||||
func WithCollector(ctx context.Context, collector *Collector) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if collector == nil {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, contextKey{}, collector)
|
||||
}
|
||||
|
||||
// FromContext returns the request timing collector, when one is active.
|
||||
func FromContext(ctx context.Context) (*Collector, bool) {
|
||||
if ctx == nil {
|
||||
return nil, false
|
||||
}
|
||||
collector, ok := ctx.Value(contextKey{}).(*Collector)
|
||||
return collector, ok && collector != nil
|
||||
}
|
||||
|
||||
// Active reports whether timing collection is enabled for this request.
|
||||
func Active(ctx context.Context) bool {
|
||||
_, ok := FromContext(ctx)
|
||||
return ok
|
||||
}
|
||||
|
||||
// Record adds a completed interval and operation count to a metric.
|
||||
func Record(ctx context.Context, name string, startedAt, endedAt time.Time, count int) {
|
||||
collector, ok := FromContext(ctx)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
collector.Record(name, startedAt, endedAt, count)
|
||||
}
|
||||
|
||||
// RecordInterval adds timing without incrementing the operation count. It is
|
||||
// useful when one logical operation has multiple blocking driver calls.
|
||||
func RecordInterval(ctx context.Context, name string, startedAt, endedAt time.Time) {
|
||||
collector, ok := FromContext(ctx)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
collector.record(name, startedAt, endedAt, 0)
|
||||
}
|
||||
|
||||
// Record adds a completed interval directly to the collector.
|
||||
func (c *Collector) Record(name string, startedAt, endedAt time.Time, count int) {
|
||||
if count <= 0 {
|
||||
count = 1
|
||||
}
|
||||
c.record(name, startedAt, endedAt, count)
|
||||
}
|
||||
|
||||
func (c *Collector) record(name string, startedAt, endedAt time.Time, count int) {
|
||||
name = normalizeMetricName(name)
|
||||
if c == nil || name == "" || startedAt.IsZero() || endedAt.Before(startedAt) {
|
||||
return
|
||||
}
|
||||
if count < 0 {
|
||||
count = 0
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
m := c.metrics[name]
|
||||
if m == nil {
|
||||
m = &metric{}
|
||||
c.metrics[name] = m
|
||||
}
|
||||
m.count += int64(count)
|
||||
if len(m.intervals) < maxIntervals {
|
||||
m.intervals = append(m.intervals, interval{start: startedAt, end: endedAt})
|
||||
}
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
// Observe starts a metric span and returns an idempotent completion function.
|
||||
func Observe(ctx context.Context, name string) func() {
|
||||
collector, ok := FromContext(ctx)
|
||||
name = normalizeMetricName(name)
|
||||
if !ok || name == "" {
|
||||
return func() {}
|
||||
}
|
||||
startedAt := time.Now()
|
||||
var once sync.Once
|
||||
return func() {
|
||||
once.Do(func() {
|
||||
collector.Record(name, startedAt, time.Now(), 1)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ObserveDependency starts a named external dependency span.
|
||||
func ObserveDependency(ctx context.Context, module string) func() {
|
||||
return Observe(ctx, dependencyMetricName(module))
|
||||
}
|
||||
|
||||
// RecordDependency records a completed external dependency interval.
|
||||
func RecordDependency(ctx context.Context, module string, startedAt, endedAt time.Time) {
|
||||
Record(ctx, dependencyMetricName(module), startedAt, endedAt, 1)
|
||||
}
|
||||
|
||||
// SetCacheStatus records the response-cache outcome for the request.
|
||||
func SetCacheStatus(ctx context.Context, status string) {
|
||||
collector, ok := FromContext(ctx)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
status = normalizeCacheStatus(status)
|
||||
if status == "" {
|
||||
return
|
||||
}
|
||||
collector.mu.Lock()
|
||||
collector.cacheStatus = status
|
||||
collector.mu.Unlock()
|
||||
}
|
||||
|
||||
// HeaderValue renders a bounded, deterministic Server-Timing header.
|
||||
func HeaderValue(ctx context.Context, endedAt time.Time, cacheStatus string) string {
|
||||
collector, ok := FromContext(ctx)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return collector.HeaderValue(endedAt, cacheStatus)
|
||||
}
|
||||
|
||||
// HeaderValue renders a bounded, deterministic Server-Timing header.
|
||||
func (c *Collector) HeaderValue(endedAt time.Time, cacheStatus string) string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
if endedAt.IsZero() {
|
||||
endedAt = time.Now()
|
||||
}
|
||||
if endedAt.Before(c.startedAt) {
|
||||
endedAt = c.startedAt
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
metrics := make(map[string]metric, len(c.metrics))
|
||||
allIntervals := make([]interval, 0)
|
||||
dependencyIntervals := make([]interval, 0)
|
||||
var dependencyCount int64
|
||||
for name, source := range c.metrics {
|
||||
copied := metric{count: source.count, intervals: append([]interval(nil), source.intervals...)}
|
||||
metrics[name] = copied
|
||||
allIntervals = append(allIntervals, copied.intervals...)
|
||||
if strings.HasPrefix(name, dependencyPrefix) {
|
||||
dependencyIntervals = append(dependencyIntervals, copied.intervals...)
|
||||
dependencyCount += copied.count
|
||||
}
|
||||
}
|
||||
storedCacheStatus := c.cacheStatus
|
||||
c.mu.Unlock()
|
||||
|
||||
total := endedAt.Sub(c.startedAt)
|
||||
blocked := unionDuration(allIntervals, c.startedAt, endedAt)
|
||||
app := total - blocked
|
||||
if app < 0 {
|
||||
app = 0
|
||||
}
|
||||
|
||||
cacheStatus = normalizeCacheStatus(cacheStatus)
|
||||
if cacheStatus == "" {
|
||||
cacheStatus = normalizeCacheStatus(storedCacheStatus)
|
||||
}
|
||||
if cacheStatus == "" {
|
||||
cacheStatus = "bypass"
|
||||
}
|
||||
|
||||
database := metrics[MetricDatabase]
|
||||
redisMetric := metrics[MetricRedis]
|
||||
parts := []string{
|
||||
"total;dur=" + formatDuration(total),
|
||||
"app;dur=" + formatDuration(app),
|
||||
fmt.Sprintf("db;dur=%s;desc=\"queries=%d\"", formatDuration(unionDuration(database.intervals, c.startedAt, endedAt)), database.count),
|
||||
fmt.Sprintf("redis;dur=%s;desc=\"commands=%d\"", formatDuration(unionDuration(redisMetric.intervals, c.startedAt, endedAt)), redisMetric.count),
|
||||
"cache;desc=\"" + cacheStatus + "\"",
|
||||
fmt.Sprintf("deps;dur=%s;desc=\"calls=%d\"", formatDuration(unionDuration(dependencyIntervals, c.startedAt, endedAt)), dependencyCount),
|
||||
}
|
||||
|
||||
dependencyNames := make([]string, 0)
|
||||
for name := range metrics {
|
||||
if strings.HasPrefix(name, dependencyPrefix) {
|
||||
dependencyNames = append(dependencyNames, name)
|
||||
}
|
||||
}
|
||||
sort.Strings(dependencyNames)
|
||||
for _, name := range dependencyNames {
|
||||
m := metrics[name]
|
||||
part := fmt.Sprintf("%s;dur=%s;desc=\"calls=%d\"", name, formatDuration(unionDuration(m.intervals, c.startedAt, endedAt)), m.count)
|
||||
candidate := strings.Join(append(parts, part), ", ")
|
||||
if len(candidate) > maxHeaderLength {
|
||||
break
|
||||
}
|
||||
parts = append(parts, part)
|
||||
}
|
||||
|
||||
return strings.Join(parts, ", ")
|
||||
}
|
||||
|
||||
func dependencyMetricName(module string) string {
|
||||
module = normalizeMetricName(module)
|
||||
module = strings.TrimPrefix(module, dependencyPrefix)
|
||||
if module == "" {
|
||||
module = "http"
|
||||
}
|
||||
return dependencyPrefix + module
|
||||
}
|
||||
|
||||
func normalizeMetricName(name string) string {
|
||||
name = strings.ToLower(strings.TrimSpace(name))
|
||||
if name == "" {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
b.Grow(min(len(name), maxMetricNameLength))
|
||||
for _, r := range name {
|
||||
if b.Len() >= maxMetricNameLength {
|
||||
break
|
||||
}
|
||||
switch {
|
||||
case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
|
||||
_, _ = b.WriteRune(r)
|
||||
case r == '_' || r == '-':
|
||||
_ = b.WriteByte('_')
|
||||
}
|
||||
}
|
||||
return strings.Trim(b.String(), "_")
|
||||
}
|
||||
|
||||
func normalizeCacheStatus(status string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(status)) {
|
||||
case "hit":
|
||||
return "hit"
|
||||
case "miss":
|
||||
return "miss"
|
||||
case "bypass":
|
||||
return "bypass"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func unionDuration(intervals []interval, lowerBound, upperBound time.Time) time.Duration {
|
||||
if len(intervals) == 0 || !upperBound.After(lowerBound) {
|
||||
return 0
|
||||
}
|
||||
normalized := make([]interval, 0, len(intervals))
|
||||
for _, item := range intervals {
|
||||
start := item.start
|
||||
end := item.end
|
||||
if start.Before(lowerBound) {
|
||||
start = lowerBound
|
||||
}
|
||||
if end.After(upperBound) {
|
||||
end = upperBound
|
||||
}
|
||||
if end.After(start) {
|
||||
normalized = append(normalized, interval{start: start, end: end})
|
||||
}
|
||||
}
|
||||
if len(normalized) == 0 {
|
||||
return 0
|
||||
}
|
||||
sort.Slice(normalized, func(i, j int) bool {
|
||||
return normalized[i].start.Before(normalized[j].start)
|
||||
})
|
||||
|
||||
currentStart := normalized[0].start
|
||||
currentEnd := normalized[0].end
|
||||
var total time.Duration
|
||||
for _, item := range normalized[1:] {
|
||||
if !item.start.After(currentEnd) {
|
||||
if item.end.After(currentEnd) {
|
||||
currentEnd = item.end
|
||||
}
|
||||
continue
|
||||
}
|
||||
total += currentEnd.Sub(currentStart)
|
||||
currentStart = item.start
|
||||
currentEnd = item.end
|
||||
}
|
||||
total += currentEnd.Sub(currentStart)
|
||||
return total
|
||||
}
|
||||
|
||||
func formatDuration(value time.Duration) string {
|
||||
if value < 0 {
|
||||
value = 0
|
||||
}
|
||||
return strconv.FormatFloat(float64(value)/float64(time.Millisecond), 'f', 1, 64)
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
package servertiming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCollectorHeaderValueAggregatesIntervals(t *testing.T) {
|
||||
startedAt := time.Unix(100, 0)
|
||||
collector := New(startedAt)
|
||||
collector.Record(MetricDatabase, startedAt.Add(10*time.Millisecond), startedAt.Add(40*time.Millisecond), 2)
|
||||
collector.Record(MetricRedis, startedAt.Add(30*time.Millisecond), startedAt.Add(50*time.Millisecond), 3)
|
||||
collector.Record(dependencyMetricName("openai"), startedAt.Add(70*time.Millisecond), startedAt.Add(100*time.Millisecond), 1)
|
||||
collector.Record(dependencyMetricName("github"), startedAt.Add(60*time.Millisecond), startedAt.Add(90*time.Millisecond), 1)
|
||||
|
||||
got := collector.HeaderValue(startedAt.Add(120*time.Millisecond), "miss")
|
||||
want := `total;dur=120.0, app;dur=40.0, db;dur=30.0;desc="queries=2", redis;dur=20.0;desc="commands=3", cache;desc="miss", deps;dur=40.0;desc="calls=2", dep_github;dur=30.0;desc="calls=1", dep_openai;dur=30.0;desc="calls=1"`
|
||||
if got != want {
|
||||
t.Fatalf("HeaderValue() = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordIntervalDoesNotIncrementCount(t *testing.T) {
|
||||
startedAt := time.Unix(200, 0)
|
||||
collector := New(startedAt)
|
||||
ctx := WithCollector(context.Background(), collector)
|
||||
|
||||
Record(ctx, MetricDatabase, startedAt.Add(10*time.Millisecond), startedAt.Add(20*time.Millisecond), 1)
|
||||
RecordInterval(ctx, MetricDatabase, startedAt.Add(30*time.Millisecond), startedAt.Add(40*time.Millisecond))
|
||||
|
||||
header := HeaderValue(ctx, startedAt.Add(100*time.Millisecond), "hit")
|
||||
if !strings.Contains(header, `db;dur=20.0;desc="queries=1"`) {
|
||||
t.Fatalf("header %q does not contain one query with both blocking intervals", header)
|
||||
}
|
||||
if !strings.Contains(header, "app;dur=80.0") {
|
||||
t.Fatalf("header %q does not subtract the interval union from app time", header)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectorCacheStatusFallback(t *testing.T) {
|
||||
startedAt := time.Unix(300, 0)
|
||||
collector := New(startedAt)
|
||||
ctx := WithCollector(context.Background(), collector)
|
||||
|
||||
SetCacheStatus(ctx, " HIT ")
|
||||
if got := HeaderValue(ctx, startedAt.Add(time.Millisecond), "invalid"); !strings.Contains(got, `cache;desc="hit"`) {
|
||||
t.Fatalf("HeaderValue() = %q, want stored cache hit", got)
|
||||
}
|
||||
|
||||
other := New(startedAt)
|
||||
if got := other.HeaderValue(startedAt.Add(time.Millisecond), "invalid"); !strings.Contains(got, `cache;desc="bypass"`) {
|
||||
t.Fatalf("HeaderValue() = %q, want cache bypass", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectorSanitizesDependencyMetric(t *testing.T) {
|
||||
startedAt := time.Unix(400, 0)
|
||||
collector := New(startedAt)
|
||||
ctx := WithCollector(context.Background(), collector)
|
||||
RecordDependency(ctx, "GitHub API\r\nInjected;dur=999", startedAt, startedAt.Add(time.Millisecond))
|
||||
|
||||
header := HeaderValue(ctx, startedAt.Add(2*time.Millisecond), "bypass")
|
||||
if strings.ContainsAny(header, "\r\n") || strings.Contains(header, ";dur=999") {
|
||||
t.Fatalf("unsafe metric content reached header: %q", header)
|
||||
}
|
||||
if !strings.Contains(header, "dep_githubapiinjecteddur999;dur=1.0") {
|
||||
t.Fatalf("sanitized dependency metric missing from header: %q", header)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectorBoundsHeaderLength(t *testing.T) {
|
||||
startedAt := time.Unix(500, 0)
|
||||
collector := New(startedAt)
|
||||
for i := 0; i < 300; i++ {
|
||||
collector.Record(
|
||||
dependencyMetricName(fmt.Sprintf("module_%03d_with_a_deliberately_long_name", i)),
|
||||
startedAt,
|
||||
startedAt.Add(time.Millisecond),
|
||||
1,
|
||||
)
|
||||
}
|
||||
|
||||
header := collector.HeaderValue(startedAt.Add(2*time.Millisecond), "bypass")
|
||||
if len(header) > maxHeaderLength {
|
||||
t.Fatalf("header length = %d, want <= %d", len(header), maxHeaderLength)
|
||||
}
|
||||
if !strings.Contains(header, "total;dur=2.0") || !strings.Contains(header, "deps;dur=1.0") {
|
||||
t.Fatalf("bounded header lost fixed metrics: %q", header)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectorConcurrentRecording(t *testing.T) {
|
||||
startedAt := time.Now()
|
||||
collector := New(startedAt)
|
||||
ctx := WithCollector(context.Background(), collector)
|
||||
|
||||
const workers = 25
|
||||
const recordsPerWorker = 100
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(workers)
|
||||
for i := 0; i < workers; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < recordsPerWorker; j++ {
|
||||
Record(ctx, MetricDatabase, startedAt, startedAt.Add(time.Microsecond), 1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
header := HeaderValue(ctx, startedAt.Add(time.Millisecond), "bypass")
|
||||
want := fmt.Sprintf(`queries=%d`, workers*recordsPerWorker)
|
||||
if !strings.Contains(header, want) {
|
||||
t.Fatalf("header %q does not contain %q", header, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestContextHelpersHandleMissingCollector(t *testing.T) {
|
||||
if Active(context.Background()) {
|
||||
t.Fatal("context without collector reported active")
|
||||
}
|
||||
if got := HeaderValue(context.Background(), time.Now(), "hit"); got != "" {
|
||||
t.Fatalf("HeaderValue() = %q without collector, want empty", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package servertiming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type dependencyModuleKey struct{}
|
||||
|
||||
type timingRoundTripper struct {
|
||||
base http.RoundTripper
|
||||
}
|
||||
|
||||
// WithDependencyModule overrides the safe module name used for an outbound call.
|
||||
func WithDependencyModule(ctx context.Context, module string) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
module = strings.TrimPrefix(normalizeMetricName(module), dependencyPrefix)
|
||||
if module == "" {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, dependencyModuleKey{}, module)
|
||||
}
|
||||
|
||||
// WrapRoundTripper records outbound response-header latency for active requests.
|
||||
func WrapRoundTripper(base http.RoundTripper) http.RoundTripper {
|
||||
if base == nil {
|
||||
base = http.DefaultTransport
|
||||
}
|
||||
if _, ok := base.(*timingRoundTripper); ok {
|
||||
return base
|
||||
}
|
||||
return &timingRoundTripper{base: base}
|
||||
}
|
||||
|
||||
// InstrumentClient returns a shallow client copy with an instrumented transport.
|
||||
func InstrumentClient(client *http.Client) *http.Client {
|
||||
if client == nil {
|
||||
client = &http.Client{}
|
||||
}
|
||||
copyClient := *client
|
||||
copyClient.Transport = WrapRoundTripper(copyClient.Transport)
|
||||
return ©Client
|
||||
}
|
||||
|
||||
// Do records response-header latency without changing the client's transport
|
||||
// type. Use it for clients whose callers inspect or configure *http.Transport.
|
||||
func Do(client *http.Client, req *http.Request) (*http.Response, error) {
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
if req == nil || !Active(req.Context()) {
|
||||
return client.Do(req)
|
||||
}
|
||||
startedAt := time.Now()
|
||||
response, err := client.Do(req)
|
||||
RecordDependency(req.Context(), dependencyModule(req), startedAt, time.Now())
|
||||
return response, err
|
||||
}
|
||||
|
||||
func (t *timingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
if req == nil || !Active(req.Context()) {
|
||||
return t.base.RoundTrip(req)
|
||||
}
|
||||
startedAt := time.Now()
|
||||
response, err := t.base.RoundTrip(req)
|
||||
RecordDependency(req.Context(), dependencyModule(req), startedAt, time.Now())
|
||||
return response, err
|
||||
}
|
||||
|
||||
func dependencyModule(req *http.Request) string {
|
||||
if req != nil {
|
||||
if module, ok := req.Context().Value(dependencyModuleKey{}).(string); ok && module != "" {
|
||||
return module
|
||||
}
|
||||
}
|
||||
if req == nil || req.URL == nil {
|
||||
return "http"
|
||||
}
|
||||
host := strings.ToLower(req.URL.Hostname())
|
||||
switch {
|
||||
case strings.Contains(host, "github"):
|
||||
return "github"
|
||||
case strings.Contains(host, "openai"):
|
||||
return "openai"
|
||||
case strings.Contains(host, "anthropic"):
|
||||
return "anthropic"
|
||||
case strings.Contains(host, "generativelanguage") || strings.Contains(host, "gemini"):
|
||||
return "gemini"
|
||||
case strings.Contains(host, "cloudcode") || strings.Contains(host, "antigravity"):
|
||||
return "antigravity"
|
||||
case strings.Contains(host, "googleapis") || strings.Contains(host, "google"):
|
||||
return "google"
|
||||
case strings.Contains(host, "amazonaws") || strings.Contains(host, "cloudflarestorage") || strings.Contains(host, "s3"):
|
||||
return "s3"
|
||||
case strings.Contains(host, "stripe") || strings.Contains(host, "airwallex") || strings.Contains(host, "alipay") || strings.Contains(host, "wechatpay") || strings.Contains(host, "paypal"):
|
||||
return "payment"
|
||||
default:
|
||||
return "http"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
package servertiming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return f(req)
|
||||
}
|
||||
|
||||
type trackingBody struct {
|
||||
read bool
|
||||
}
|
||||
|
||||
func (b *trackingBody) Read(_ []byte) (int, error) {
|
||||
b.read = true
|
||||
return 0, io.EOF
|
||||
}
|
||||
|
||||
func (b *trackingBody) Close() error { return nil }
|
||||
|
||||
func TestWrapRoundTripperRecordsResponseHeaderLatency(t *testing.T) {
|
||||
startedAt := time.Now()
|
||||
collector := New(startedAt)
|
||||
body := &trackingBody{}
|
||||
baseCalled := false
|
||||
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
baseCalled = true
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: body,
|
||||
Header: make(http.Header),
|
||||
Request: req,
|
||||
}, nil
|
||||
})
|
||||
req, err := http.NewRequestWithContext(WithCollector(context.Background(), collector), http.MethodGet, "https://api.github.com/repos/example/project", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
resp, err := WrapRoundTripper(base).RoundTrip(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if !baseCalled {
|
||||
t.Fatal("base RoundTripper was not called")
|
||||
}
|
||||
if body.read {
|
||||
t.Fatal("RoundTripper instrumentation read the response body; timing must stop at response headers")
|
||||
}
|
||||
header := collector.HeaderValue(time.Now(), "bypass")
|
||||
if !strings.Contains(header, `dep_github;dur=`) || !strings.Contains(header, `deps;dur=`) {
|
||||
t.Fatalf("dependency metrics missing from header: %q", header)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWrapRoundTripperUsesContextModuleOverride(t *testing.T) {
|
||||
collector := New(time.Now())
|
||||
ctx := WithDependencyModule(WithCollector(context.Background(), collector), "data-managementd")
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://private.example.test/path", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
|
||||
})
|
||||
|
||||
if _, err := WrapRoundTripper(base).RoundTrip(req); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
header := collector.HeaderValue(time.Now(), "bypass")
|
||||
if !strings.Contains(header, "dep_data_managementd") {
|
||||
t.Fatalf("module override missing from header: %q", header)
|
||||
}
|
||||
if strings.Contains(header, "private.example") {
|
||||
t.Fatalf("raw host leaked into header: %q", header)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWrapRoundTripperSkipsInactiveContext(t *testing.T) {
|
||||
called := false
|
||||
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
called = true
|
||||
return &http.Response{StatusCode: http.StatusNoContent, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
|
||||
})
|
||||
req, err := http.NewRequest(http.MethodGet, "https://api.openai.com/v1/models", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := WrapRoundTripper(base).RoundTrip(req); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !called {
|
||||
t.Fatal("inactive request did not reach base RoundTripper")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDoRecordsWithoutChangingTransportType(t *testing.T) {
|
||||
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
|
||||
})
|
||||
client := &http.Client{Transport: base}
|
||||
collector := New(time.Now())
|
||||
req, err := http.NewRequestWithContext(WithCollector(context.Background(), collector), http.MethodGet, "https://api.openai.com/v1/models", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := Do(client, req); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, ok := client.Transport.(roundTripFunc); !ok {
|
||||
t.Fatalf("Do changed client transport type to %T", client.Transport)
|
||||
}
|
||||
if header := collector.HeaderValue(time.Now(), "bypass"); !strings.Contains(header, "dep_openai;dur=") {
|
||||
t.Fatalf("dependency metric missing from header: %q", header)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDependencyModuleClassification(t *testing.T) {
|
||||
tests := map[string]string{
|
||||
"https://api.github.com/repos/a/b": "github",
|
||||
"https://api.openai.com/v1/models": "openai",
|
||||
"https://api.anthropic.com/v1/messages": "anthropic",
|
||||
"https://generativelanguage.googleapis.com/v1/models": "gemini",
|
||||
"https://cloudcode-pa.googleapis.com/v1internal": "antigravity",
|
||||
"https://storage.googleapis.com/bucket/object": "google",
|
||||
"https://bucket.s3.amazonaws.com/object": "s3",
|
||||
"https://api.stripe.com/v1/refunds": "payment",
|
||||
"https://dependency.example.test/path": "http",
|
||||
}
|
||||
for rawURL, want := range tests {
|
||||
req, err := http.NewRequest(http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewRequest(%q): %v", rawURL, err)
|
||||
}
|
||||
if got := dependencyModule(req); got != want {
|
||||
t.Errorf("dependencyModule(%q) = %q, want %q", rawURL, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientInstrumentationDoesNotMutateOriginal(t *testing.T) {
|
||||
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
|
||||
})
|
||||
original := &http.Client{Transport: base, Timeout: time.Second}
|
||||
instrumented := InstrumentClient(original)
|
||||
if instrumented == original {
|
||||
t.Fatal("InstrumentClient returned the original client")
|
||||
}
|
||||
if _, ok := original.Transport.(roundTripFunc); !ok {
|
||||
t.Fatalf("InstrumentClient mutated the original transport to %T", original.Transport)
|
||||
}
|
||||
if instrumented.Timeout != original.Timeout {
|
||||
t.Fatal("InstrumentClient did not preserve client settings")
|
||||
}
|
||||
if WrapRoundTripper(instrumented.Transport) != instrumented.Transport {
|
||||
t.Fatal("WrapRoundTripper wrapped an already instrumented transport twice")
|
||||
}
|
||||
}
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
@@ -63,12 +64,14 @@ func (s *S3BackupStore) Upload(ctx context.Context, key string, body io.Reader,
|
||||
return 0, fmt.Errorf("read body: %w", err)
|
||||
}
|
||||
|
||||
finish := servertiming.ObserveDependency(ctx, "s3")
|
||||
_, err = s.client.PutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: &s.bucket,
|
||||
Key: &key,
|
||||
Body: bytes.NewReader(data),
|
||||
ContentType: &contentType,
|
||||
})
|
||||
finish()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("S3 PutObject: %w", err)
|
||||
}
|
||||
@@ -76,10 +79,12 @@ func (s *S3BackupStore) Upload(ctx context.Context, key string, body io.Reader,
|
||||
}
|
||||
|
||||
func (s *S3BackupStore) Download(ctx context.Context, key string) (io.ReadCloser, error) {
|
||||
finish := servertiming.ObserveDependency(ctx, "s3")
|
||||
result, err := s.client.GetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: &s.bucket,
|
||||
Key: &key,
|
||||
})
|
||||
finish()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("S3 GetObject: %w", err)
|
||||
}
|
||||
@@ -87,10 +92,12 @@ func (s *S3BackupStore) Download(ctx context.Context, key string) (io.ReadCloser
|
||||
}
|
||||
|
||||
func (s *S3BackupStore) Delete(ctx context.Context, key string) error {
|
||||
finish := servertiming.ObserveDependency(ctx, "s3")
|
||||
_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
|
||||
Bucket: &s.bucket,
|
||||
Key: &key,
|
||||
})
|
||||
finish()
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -107,9 +114,11 @@ func (s *S3BackupStore) PresignURL(ctx context.Context, key string, expiry time.
|
||||
}
|
||||
|
||||
func (s *S3BackupStore) HeadBucket(ctx context.Context) error {
|
||||
finish := servertiming.ObserveDependency(ctx, "s3")
|
||||
_, err := s.client.HeadBucket(ctx, &s3.HeadBucketInput{
|
||||
Bucket: &s.bucket,
|
||||
})
|
||||
finish()
|
||||
if err != nil {
|
||||
return fmt.Errorf("S3 HeadBucket failed: %w", err)
|
||||
}
|
||||
|
||||
@@ -276,5 +276,5 @@ func createReqClient(proxyURL string) (*req.Client, error) {
|
||||
client.SetProxyURL(trimmed)
|
||||
}
|
||||
|
||||
return client, nil
|
||||
return instrumentReqClient(client), nil
|
||||
}
|
||||
|
||||
@@ -15,7 +15,7 @@ import (
|
||||
|
||||
"entgo.io/ent/dialect"
|
||||
entsql "entgo.io/ent/dialect/sql"
|
||||
_ "github.com/lib/pq" // PostgreSQL 驱动,通过副作用导入注册驱动
|
||||
"github.com/lib/pq"
|
||||
)
|
||||
|
||||
// InitEnt 初始化 Ent ORM 客户端并返回客户端实例和底层的 *sql.DB。
|
||||
@@ -48,9 +48,19 @@ func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) {
|
||||
|
||||
// 使用 Ent 的 SQL 驱动打开 PostgreSQL 连接。
|
||||
// dialect.Postgres 指定使用 PostgreSQL 方言进行 SQL 生成。
|
||||
drv, err := entsql.Open(dialect.Postgres, dsn)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
var drv *entsql.Driver
|
||||
if cfg.Server.EnableServerTiming {
|
||||
connector, err := pq.NewConnector(dsn)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
drv = entsql.OpenDB(dialect.Postgres, sql.OpenDB(newServerTimingConnector(connector)))
|
||||
} else {
|
||||
var err error
|
||||
drv, err = entsql.Open(dialect.Postgres, dsn)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
}
|
||||
applyDBPoolSettings(drv.DB(), cfg)
|
||||
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
|
||||
@@ -186,7 +187,7 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i
|
||||
}
|
||||
|
||||
// 执行请求
|
||||
resp, err := entry.client.Do(req)
|
||||
resp, err := servertiming.Do(entry.client, req)
|
||||
if err != nil {
|
||||
s.recordOpenAIHTTP2Failure(profile, entry.protocolMode, entry.proxyKey, err)
|
||||
// 请求失败,立即减少计数
|
||||
@@ -243,7 +244,7 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := entry.client.Do(req)
|
||||
resp, err := servertiming.Do(entry.client, req)
|
||||
if err != nil {
|
||||
atomic.AddInt64(&entry.inFlight, -1)
|
||||
atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
|
||||
|
||||
@@ -21,7 +21,11 @@ import (
|
||||
// 2. MinIdleConns: 保持最小空闲连接,减少冷启动延迟(默认 10)
|
||||
// 3. DialTimeout/ReadTimeout/WriteTimeout: 精确控制各阶段超时
|
||||
func InitRedis(cfg *config.Config) *redis.Client {
|
||||
return redis.NewClient(buildRedisOptions(cfg))
|
||||
client := redis.NewClient(buildRedisOptions(cfg))
|
||||
if cfg.Server.EnableServerTiming {
|
||||
client.AddHook(serverTimingRedisHook{})
|
||||
}
|
||||
return client
|
||||
}
|
||||
|
||||
// buildRedisOptions 构建 Redis 连接选项
|
||||
|
||||
@@ -2,11 +2,13 @@ package repository
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
|
||||
"github.com/imroc/req/v3"
|
||||
)
|
||||
@@ -57,6 +59,7 @@ func getSharedReqClient(opts reqClientOptions) (*req.Client, error) {
|
||||
if trimmed != "" {
|
||||
client.SetProxyURL(trimmed)
|
||||
}
|
||||
client = instrumentReqClient(client)
|
||||
|
||||
actual, _ := sharedReqClients.LoadOrStore(key, client)
|
||||
if c, ok := actual.(*req.Client); ok {
|
||||
@@ -65,6 +68,17 @@ func getSharedReqClient(opts reqClientOptions) (*req.Client, error) {
|
||||
return client, nil
|
||||
}
|
||||
|
||||
func instrumentReqClient(client *req.Client) *req.Client {
|
||||
if client == nil {
|
||||
return nil
|
||||
}
|
||||
client.GetTransport().WrapRoundTripFunc(func(rt http.RoundTripper) req.HttpRoundTripFunc {
|
||||
timed := servertiming.WrapRoundTripper(rt)
|
||||
return timed.RoundTrip
|
||||
})
|
||||
return client
|
||||
}
|
||||
|
||||
func buildReqClientKey(opts reqClientOptions) string {
|
||||
return fmt.Sprintf("%s|%s|%t|%t",
|
||||
strings.TrimSpace(opts.ProxyURL),
|
||||
|
||||
@@ -1,12 +1,17 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
"unsafe"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/imroc/req/v3"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
@@ -118,3 +123,20 @@ func TestCreateGeminiReqClient_ForceHTTP2Disabled(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "", forceHTTPVersion(t, client))
|
||||
}
|
||||
|
||||
func TestInstrumentReqClientRecordsDependency(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
collector := servertiming.New(time.Now())
|
||||
ctx := servertiming.WithCollector(context.Background(), collector)
|
||||
client := instrumentReqClient(req.C())
|
||||
response, err := client.R().SetContext(ctx).Get(server.URL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusNoContent, response.StatusCode)
|
||||
|
||||
header := collector.HeaderValue(time.Now(), "bypass")
|
||||
require.True(t, strings.Contains(header, "dep_http;dur="), header)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
type serverTimingRedisHook struct{}
|
||||
|
||||
func (serverTimingRedisHook) DialHook(next redis.DialHook) redis.DialHook {
|
||||
return next
|
||||
}
|
||||
|
||||
func (serverTimingRedisHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook {
|
||||
return func(ctx context.Context, cmd redis.Cmder) error {
|
||||
if !servertiming.Active(ctx) {
|
||||
return next(ctx, cmd)
|
||||
}
|
||||
startedAt := time.Now()
|
||||
err := next(ctx, cmd)
|
||||
servertiming.Record(ctx, servertiming.MetricRedis, startedAt, time.Now(), 1)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
func (serverTimingRedisHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook {
|
||||
return func(ctx context.Context, cmds []redis.Cmder) error {
|
||||
if !servertiming.Active(ctx) {
|
||||
return next(ctx, cmds)
|
||||
}
|
||||
startedAt := time.Now()
|
||||
err := next(ctx, cmds)
|
||||
servertiming.Record(ctx, servertiming.MetricRedis, startedAt, time.Now(), len(cmds))
|
||||
return err
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
func TestServerTimingRedisHookRecordsCommands(t *testing.T) {
|
||||
collector := servertiming.New(time.Now())
|
||||
ctx := servertiming.WithCollector(context.Background(), collector)
|
||||
hook := serverTimingRedisHook{}
|
||||
|
||||
process := hook.ProcessHook(func(context.Context, redis.Cmder) error {
|
||||
time.Sleep(time.Millisecond)
|
||||
return errors.New("redis failure")
|
||||
})
|
||||
if err := process(ctx, redis.NewStringCmd(ctx, "get", "sensitive-key")); err == nil {
|
||||
t.Fatal("ProcessHook did not return the underlying error")
|
||||
}
|
||||
|
||||
pipeline := hook.ProcessPipelineHook(func(context.Context, []redis.Cmder) error {
|
||||
time.Sleep(time.Millisecond)
|
||||
return nil
|
||||
})
|
||||
commands := []redis.Cmder{
|
||||
redis.NewStringCmd(ctx, "get", "first-secret"),
|
||||
redis.NewStringCmd(ctx, "get", "second-secret"),
|
||||
redis.NewStatusCmd(ctx, "set", "third-secret", "value"),
|
||||
}
|
||||
if err := pipeline(ctx, commands); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
header := collector.HeaderValue(time.Now(), "bypass")
|
||||
if !strings.Contains(header, `commands=4`) {
|
||||
t.Fatalf("header %q does not report one command and a three-command pipeline", header)
|
||||
}
|
||||
if strings.Contains(header, "secret") || strings.Contains(header, "get") {
|
||||
t.Fatalf("Redis command details leaked into header: %q", header)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingRedisHookSkipsInactiveContext(t *testing.T) {
|
||||
called := false
|
||||
hook := serverTimingRedisHook{}
|
||||
process := hook.ProcessHook(func(context.Context, redis.Cmder) error {
|
||||
called = true
|
||||
return nil
|
||||
})
|
||||
ctx := context.Background()
|
||||
if err := process(ctx, redis.NewStringCmd(ctx, "ping")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !called {
|
||||
t.Fatal("inactive Redis command did not reach the next hook")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql/driver"
|
||||
"errors"
|
||||
"io"
|
||||
"reflect"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
)
|
||||
|
||||
type serverTimingConnector struct {
|
||||
base driver.Connector
|
||||
}
|
||||
|
||||
func newServerTimingConnector(base driver.Connector) driver.Connector {
|
||||
return &serverTimingConnector{base: base}
|
||||
}
|
||||
|
||||
func (c *serverTimingConnector) Connect(ctx context.Context) (driver.Conn, error) {
|
||||
startedAt := time.Now()
|
||||
conn, err := c.base.Connect(ctx)
|
||||
servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &serverTimingConn{Conn: conn}, nil
|
||||
}
|
||||
|
||||
func (c *serverTimingConnector) Driver() driver.Driver {
|
||||
return c.base.Driver()
|
||||
}
|
||||
|
||||
type serverTimingConn struct {
|
||||
driver.Conn
|
||||
}
|
||||
|
||||
func (c *serverTimingConn) Prepare(query string) (driver.Stmt, error) {
|
||||
stmt, err := c.Conn.Prepare(query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &serverTimingStmt{Stmt: stmt}, nil
|
||||
}
|
||||
|
||||
func (c *serverTimingConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
|
||||
startedAt := time.Now()
|
||||
var (
|
||||
stmt driver.Stmt
|
||||
err error
|
||||
)
|
||||
if preparer, ok := c.Conn.(driver.ConnPrepareContext); ok {
|
||||
stmt, err = preparer.PrepareContext(ctx, query)
|
||||
} else {
|
||||
stmt, err = c.Conn.Prepare(query)
|
||||
}
|
||||
servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &serverTimingStmt{Stmt: stmt}, nil
|
||||
}
|
||||
|
||||
func (c *serverTimingConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
|
||||
execer, ok := c.Conn.(driver.ExecerContext)
|
||||
if !ok {
|
||||
return nil, driver.ErrSkip
|
||||
}
|
||||
startedAt := time.Now()
|
||||
result, err := execer.ExecContext(ctx, query, args)
|
||||
servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (c *serverTimingConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
|
||||
queryer, ok := c.Conn.(driver.QueryerContext)
|
||||
if !ok {
|
||||
return nil, driver.ErrSkip
|
||||
}
|
||||
startedAt := time.Now()
|
||||
rows, err := queryer.QueryContext(ctx, query, args)
|
||||
servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
|
||||
if err != nil || rows == nil {
|
||||
return rows, err
|
||||
}
|
||||
return newServerTimingRows(ctx, rows), nil
|
||||
}
|
||||
|
||||
func (c *serverTimingConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
|
||||
startedAt := time.Now()
|
||||
var (
|
||||
tx driver.Tx
|
||||
err error
|
||||
)
|
||||
if beginner, ok := c.Conn.(driver.ConnBeginTx); ok {
|
||||
tx, err = beginner.BeginTx(ctx, opts)
|
||||
} else {
|
||||
if opts.Isolation != driver.IsolationLevel(0) {
|
||||
return nil, errors.New("driver does not support non-default isolation")
|
||||
}
|
||||
if opts.ReadOnly {
|
||||
return nil, errors.New("driver does not support read-only transactions")
|
||||
}
|
||||
// The wrapper exposes ConnBeginTx, so it must retain database/sql's
|
||||
// legacy fallback for drivers that only implement Conn.Begin.
|
||||
tx, err = c.Conn.Begin() //nolint:staticcheck // Required driver compatibility fallback.
|
||||
}
|
||||
servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
|
||||
if err != nil || tx == nil {
|
||||
return tx, err
|
||||
}
|
||||
return &serverTimingTx{Tx: tx, ctx: ctx}, nil
|
||||
}
|
||||
|
||||
func (c *serverTimingConn) Ping(ctx context.Context) error {
|
||||
if pinger, ok := c.Conn.(driver.Pinger); ok {
|
||||
startedAt := time.Now()
|
||||
err := pinger.Ping(ctx)
|
||||
servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *serverTimingConn) ResetSession(ctx context.Context) error {
|
||||
if resetter, ok := c.Conn.(driver.SessionResetter); ok {
|
||||
startedAt := time.Now()
|
||||
err := resetter.ResetSession(ctx)
|
||||
servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *serverTimingConn) IsValid() bool {
|
||||
if validator, ok := c.Conn.(driver.Validator); ok {
|
||||
return validator.IsValid()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (c *serverTimingConn) CheckNamedValue(value *driver.NamedValue) error {
|
||||
if checker, ok := c.Conn.(driver.NamedValueChecker); ok {
|
||||
return checker.CheckNamedValue(value)
|
||||
}
|
||||
return driver.ErrSkip
|
||||
}
|
||||
|
||||
type serverTimingStmt struct {
|
||||
driver.Stmt
|
||||
}
|
||||
|
||||
func (s *serverTimingStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
|
||||
startedAt := time.Now()
|
||||
var (
|
||||
result driver.Result
|
||||
err error
|
||||
)
|
||||
if execer, ok := s.Stmt.(driver.StmtExecContext); ok {
|
||||
result, err = execer.ExecContext(ctx, args)
|
||||
} else {
|
||||
var values []driver.Value
|
||||
values, err = namedValues(args)
|
||||
if err == nil {
|
||||
// The wrapper exposes StmtExecContext and must preserve the fallback
|
||||
// database/sql would use for a legacy driver statement.
|
||||
result, err = s.Stmt.Exec(values) //nolint:staticcheck // Required driver compatibility fallback.
|
||||
}
|
||||
}
|
||||
servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (s *serverTimingStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
|
||||
startedAt := time.Now()
|
||||
var (
|
||||
rows driver.Rows
|
||||
err error
|
||||
)
|
||||
if queryer, ok := s.Stmt.(driver.StmtQueryContext); ok {
|
||||
rows, err = queryer.QueryContext(ctx, args)
|
||||
} else {
|
||||
var values []driver.Value
|
||||
values, err = namedValues(args)
|
||||
if err == nil {
|
||||
// The wrapper exposes StmtQueryContext and must preserve the fallback
|
||||
// database/sql would use for a legacy driver statement.
|
||||
rows, err = s.Stmt.Query(values) //nolint:staticcheck // Required driver compatibility fallback.
|
||||
}
|
||||
}
|
||||
servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
|
||||
if err != nil || rows == nil {
|
||||
return rows, err
|
||||
}
|
||||
return newServerTimingRows(ctx, rows), nil
|
||||
}
|
||||
|
||||
func (s *serverTimingStmt) CheckNamedValue(value *driver.NamedValue) error {
|
||||
if checker, ok := s.Stmt.(driver.NamedValueChecker); ok {
|
||||
return checker.CheckNamedValue(value)
|
||||
}
|
||||
return driver.ErrSkip
|
||||
}
|
||||
|
||||
func namedValues(args []driver.NamedValue) ([]driver.Value, error) {
|
||||
values := make([]driver.Value, len(args))
|
||||
for i, arg := range args {
|
||||
if arg.Name != "" {
|
||||
return nil, errors.New("named parameters are not supported")
|
||||
}
|
||||
values[i] = arg.Value
|
||||
}
|
||||
return values, nil
|
||||
}
|
||||
|
||||
type serverTimingRows struct {
|
||||
driver.Rows
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func newServerTimingRows(ctx context.Context, rows driver.Rows) *serverTimingRows {
|
||||
return &serverTimingRows{Rows: rows, ctx: ctx}
|
||||
}
|
||||
|
||||
func (r *serverTimingRows) Close() error {
|
||||
startedAt := time.Now()
|
||||
err := r.Rows.Close()
|
||||
servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now())
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *serverTimingRows) Next(dest []driver.Value) error {
|
||||
startedAt := time.Now()
|
||||
err := r.Rows.Next(dest)
|
||||
servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now())
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *serverTimingRows) HasNextResultSet() bool {
|
||||
if rows, ok := r.Rows.(driver.RowsNextResultSet); ok {
|
||||
return rows.HasNextResultSet()
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (r *serverTimingRows) NextResultSet() error {
|
||||
rows, ok := r.Rows.(driver.RowsNextResultSet)
|
||||
if !ok {
|
||||
return io.EOF
|
||||
}
|
||||
startedAt := time.Now()
|
||||
err := rows.NextResultSet()
|
||||
servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now())
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *serverTimingRows) ColumnTypeScanType(index int) reflect.Type {
|
||||
if rows, ok := r.Rows.(driver.RowsColumnTypeScanType); ok {
|
||||
return rows.ColumnTypeScanType(index)
|
||||
}
|
||||
return reflect.TypeOf(new(any)).Elem()
|
||||
}
|
||||
|
||||
func (r *serverTimingRows) ColumnTypeDatabaseTypeName(index int) string {
|
||||
if rows, ok := r.Rows.(driver.RowsColumnTypeDatabaseTypeName); ok {
|
||||
return rows.ColumnTypeDatabaseTypeName(index)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (r *serverTimingRows) ColumnTypeLength(index int) (int64, bool) {
|
||||
if rows, ok := r.Rows.(driver.RowsColumnTypeLength); ok {
|
||||
return rows.ColumnTypeLength(index)
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func (r *serverTimingRows) ColumnTypeNullable(index int) (bool, bool) {
|
||||
if rows, ok := r.Rows.(driver.RowsColumnTypeNullable); ok {
|
||||
return rows.ColumnTypeNullable(index)
|
||||
}
|
||||
return false, false
|
||||
}
|
||||
|
||||
func (r *serverTimingRows) ColumnTypePrecisionScale(index int) (int64, int64, bool) {
|
||||
if rows, ok := r.Rows.(driver.RowsColumnTypePrecisionScale); ok {
|
||||
return rows.ColumnTypePrecisionScale(index)
|
||||
}
|
||||
return 0, 0, false
|
||||
}
|
||||
|
||||
type serverTimingTx struct {
|
||||
driver.Tx
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func (t *serverTimingTx) Commit() error {
|
||||
startedAt := time.Now()
|
||||
err := t.Tx.Commit()
|
||||
servertiming.RecordInterval(t.ctx, servertiming.MetricDatabase, startedAt, time.Now())
|
||||
return err
|
||||
}
|
||||
|
||||
func (t *serverTimingTx) Rollback() error {
|
||||
startedAt := time.Now()
|
||||
err := t.Tx.Rollback()
|
||||
servertiming.RecordInterval(t.ctx, servertiming.MetricDatabase, startedAt, time.Now())
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql/driver"
|
||||
"io"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
)
|
||||
|
||||
const fakeDriverDelay = 2 * time.Millisecond
|
||||
|
||||
type timingFakeDriver struct{}
|
||||
|
||||
func (timingFakeDriver) Open(string) (driver.Conn, error) { return newTimingFakeConn(), nil }
|
||||
|
||||
type timingFakeConnector struct {
|
||||
conn driver.Conn
|
||||
}
|
||||
|
||||
func (c timingFakeConnector) Connect(context.Context) (driver.Conn, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return c.conn, nil
|
||||
}
|
||||
|
||||
func (timingFakeConnector) Driver() driver.Driver { return timingFakeDriver{} }
|
||||
|
||||
type timingFakeConn struct{}
|
||||
|
||||
func newTimingFakeConn() *timingFakeConn { return &timingFakeConn{} }
|
||||
|
||||
func (c *timingFakeConn) Prepare(string) (driver.Stmt, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return &timingFakeStmt{}, nil
|
||||
}
|
||||
|
||||
func (c *timingFakeConn) PrepareContext(context.Context, string) (driver.Stmt, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return &timingFakeStmt{}, nil
|
||||
}
|
||||
|
||||
func (c *timingFakeConn) Close() error { return nil }
|
||||
|
||||
func (c *timingFakeConn) Begin() (driver.Tx, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return &timingFakeTx{}, nil
|
||||
}
|
||||
|
||||
func (c *timingFakeConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return &timingFakeTx{}, nil
|
||||
}
|
||||
|
||||
func (c *timingFakeConn) ExecContext(context.Context, string, []driver.NamedValue) (driver.Result, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return driver.RowsAffected(1), nil
|
||||
}
|
||||
|
||||
func (c *timingFakeConn) QueryContext(context.Context, string, []driver.NamedValue) (driver.Rows, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil
|
||||
}
|
||||
|
||||
func (c *timingFakeConn) Ping(context.Context) error {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *timingFakeConn) ResetSession(context.Context) error {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return nil
|
||||
}
|
||||
|
||||
type timingFakeStmt struct{}
|
||||
|
||||
func (s *timingFakeStmt) Close() error { return nil }
|
||||
func (s *timingFakeStmt) NumInput() int { return -1 }
|
||||
|
||||
func (s *timingFakeStmt) Exec([]driver.Value) (driver.Result, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return driver.RowsAffected(1), nil
|
||||
}
|
||||
|
||||
func (s *timingFakeStmt) Query([]driver.Value) (driver.Rows, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil
|
||||
}
|
||||
|
||||
func (s *timingFakeStmt) ExecContext(context.Context, []driver.NamedValue) (driver.Result, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return driver.RowsAffected(1), nil
|
||||
}
|
||||
|
||||
func (s *timingFakeStmt) QueryContext(context.Context, []driver.NamedValue) (driver.Rows, error) {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil
|
||||
}
|
||||
|
||||
type timingFakeRows struct {
|
||||
values [][]driver.Value
|
||||
index int
|
||||
}
|
||||
|
||||
func (r *timingFakeRows) Columns() []string { return []string{"value"} }
|
||||
|
||||
func (r *timingFakeRows) Close() error {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *timingFakeRows) Next(dest []driver.Value) error {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
if r.index >= len(r.values) {
|
||||
return io.EOF
|
||||
}
|
||||
copy(dest, r.values[r.index])
|
||||
r.index++
|
||||
return nil
|
||||
}
|
||||
|
||||
type timingFakeTx struct{}
|
||||
|
||||
func (t *timingFakeTx) Commit() error {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *timingFakeTx) Rollback() error {
|
||||
time.Sleep(fakeDriverDelay)
|
||||
return nil
|
||||
}
|
||||
|
||||
func metricDuration(t *testing.T, header, metric string) float64 {
|
||||
t.Helper()
|
||||
re := regexp.MustCompile(`(?:^|, )` + regexp.QuoteMeta(metric) + `;dur=([0-9]+(?:\.[0-9]+)?)`)
|
||||
match := re.FindStringSubmatch(header)
|
||||
if len(match) != 2 {
|
||||
t.Fatalf("metric %q missing from header %q", metric, header)
|
||||
}
|
||||
value, err := strconv.ParseFloat(match[1], 64)
|
||||
if err != nil {
|
||||
t.Fatalf("parse %s duration: %v", metric, err)
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func TestServerTimingConnectorRecordsDriverCallsWithoutRowLifetime(t *testing.T) {
|
||||
startedAt := time.Now()
|
||||
collector := servertiming.New(startedAt)
|
||||
ctx := servertiming.WithCollector(context.Background(), collector)
|
||||
|
||||
wrapped := newServerTimingConnector(timingFakeConnector{conn: newTimingFakeConn()})
|
||||
rawConn, err := wrapped.Connect(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
conn, ok := rawConn.(*serverTimingConn)
|
||||
if !ok {
|
||||
t.Fatalf("Connect() returned %T, want *serverTimingConn", rawConn)
|
||||
}
|
||||
|
||||
if _, err := conn.ExecContext(ctx, "sensitive update", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rows, err := conn.QueryContext(ctx, "sensitive select", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
values := make([]driver.Value, 1)
|
||||
if err := rows.Next(values); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Application work between row reads must remain app time.
|
||||
time.Sleep(30 * time.Millisecond)
|
||||
if err := rows.Next(values); err != io.EOF {
|
||||
t.Fatalf("rows.Next() = %v, want EOF", err)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
header := collector.HeaderValue(time.Now(), "bypass")
|
||||
if !strings.Contains(header, `queries=2`) {
|
||||
t.Fatalf("header %q does not report two SQL operations", header)
|
||||
}
|
||||
if strings.Contains(header, "sensitive") {
|
||||
t.Fatalf("SQL text leaked into header: %q", header)
|
||||
}
|
||||
if app, db := metricDuration(t, header, "app"), metricDuration(t, header, "db"); app <= db {
|
||||
t.Fatalf("row processing gap was counted as DB time: app=%.1fms db=%.1fms header=%q", app, db, header)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingPreparedStatementsAndTransactions(t *testing.T) {
|
||||
collector := servertiming.New(time.Now())
|
||||
ctx := servertiming.WithCollector(context.Background(), collector)
|
||||
conn := &serverTimingConn{Conn: newTimingFakeConn()}
|
||||
|
||||
stmt, err := conn.PrepareContext(ctx, "prepare sensitive statement")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
timedStmt, ok := stmt.(*serverTimingStmt)
|
||||
if !ok {
|
||||
t.Fatalf("PrepareContext() returned %T, want *serverTimingStmt", stmt)
|
||||
}
|
||||
if _, err := timedStmt.ExecContext(ctx, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rows, err := timedStmt.QueryContext(ctx, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
tx, err := conn.BeginTx(ctx, driver.TxOptions{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := conn.Ping(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := conn.ResetSession(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
header := collector.HeaderValue(time.Now(), "bypass")
|
||||
if !strings.Contains(header, `queries=3`) {
|
||||
t.Fatalf("header %q does not report prepare, exec, and query operations", header)
|
||||
}
|
||||
if metricDuration(t, header, "db") <= 0 {
|
||||
t.Fatalf("DB duration was not recorded: %q", header)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNamedValuesRejectNamedParameters(t *testing.T) {
|
||||
if _, err := namedValues([]driver.NamedValue{{Name: "secret", Value: 1}}); err == nil {
|
||||
t.Fatal("namedValues accepted a named parameter")
|
||||
}
|
||||
values, err := namedValues([]driver.NamedValue{{Ordinal: 1, Value: "value"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(values) != 1 || values[0] != "value" {
|
||||
t.Fatalf("namedValues() = %#v", values)
|
||||
}
|
||||
}
|
||||
@@ -52,7 +52,7 @@ func CORS(cfg config.CORSConfig) gin.HandlerFunc {
|
||||
}
|
||||
allowHeaders := []string{
|
||||
"Content-Type", "Content-Length", "Accept-Encoding", "X-CSRF-Token", "Authorization",
|
||||
"accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key",
|
||||
"accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key", "X-Admin-UI-Request",
|
||||
}
|
||||
// OpenAI Node SDK 会发送 x-stainless-* 请求头,需在 CORS 中显式放行。
|
||||
openAIProperties := []string{
|
||||
@@ -83,7 +83,7 @@ func CORS(cfg config.CORSConfig) gin.HandlerFunc {
|
||||
}
|
||||
c.Writer.Header().Set("Access-Control-Allow-Headers", allowHeadersValue)
|
||||
c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE, PATCH")
|
||||
c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag")
|
||||
c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag, Server-Timing")
|
||||
c.Writer.Header().Set("Access-Control-Max-Age", "86400")
|
||||
}
|
||||
// 处理预检请求
|
||||
|
||||
@@ -103,8 +103,10 @@ func TestCORS_AllowedOrigin_HasAllowHeaders(t *testing.T) {
|
||||
// 应设置 Allow-Headers、Allow-Methods 和 Max-Age
|
||||
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"),
|
||||
"允许的 origin 应收到 Allow-Headers")
|
||||
assert.Contains(t, w.Header().Get("Access-Control-Allow-Headers"), "X-Admin-UI-Request")
|
||||
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Methods"),
|
||||
"允许的 origin 应收到 Allow-Methods")
|
||||
assert.Contains(t, w.Header().Get("Access-Control-Expose-Headers"), "Server-Timing")
|
||||
assert.Equal(t, "86400", w.Header().Get("Access-Control-Max-Age"),
|
||||
"允许的 origin 应收到 Max-Age=86400")
|
||||
assert.Equal(t, "https://allowed.example.com", w.Header().Get("Access-Control-Allow-Origin"),
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
snapshotCacheHeader = "X-Snapshot-Cache"
|
||||
usageCacheHeader = "X-Usage-Stats-Cache"
|
||||
)
|
||||
|
||||
type serverTimingResponseWriter struct {
|
||||
gin.ResponseWriter
|
||||
context *gin.Context
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) Unwrap() http.ResponseWriter {
|
||||
return w.ResponseWriter
|
||||
}
|
||||
|
||||
// ServerTiming collects timing only for requests made by the Admin web UI.
|
||||
func ServerTiming(enabled bool) gin.HandlerFunc {
|
||||
if !enabled {
|
||||
return func(c *gin.Context) {
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
return func(c *gin.Context) {
|
||||
if !isAdminUIRequest(c) || c.Request == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
collector := servertiming.New(time.Now())
|
||||
c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector))
|
||||
writer := &serverTimingResponseWriter{
|
||||
ResponseWriter: c.Writer,
|
||||
context: c,
|
||||
}
|
||||
c.Writer = writer
|
||||
c.Next()
|
||||
writer.finalize()
|
||||
}
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) WriteHeader(statusCode int) {
|
||||
w.ResponseWriter.WriteHeader(statusCode)
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) WriteHeaderNow() {
|
||||
w.finalize()
|
||||
w.ResponseWriter.WriteHeaderNow()
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) Write(data []byte) (int, error) {
|
||||
w.finalize()
|
||||
return w.ResponseWriter.Write(data)
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) WriteString(data string) (int, error) {
|
||||
w.finalize()
|
||||
return w.ResponseWriter.WriteString(data)
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) Flush() {
|
||||
w.finalize()
|
||||
w.ResponseWriter.Flush()
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) finalize() {
|
||||
if w == nil {
|
||||
return
|
||||
}
|
||||
w.once.Do(func() {
|
||||
if value := ServerTimingHeaderValue(w.context); value != "" {
|
||||
w.ResponseWriter.Header().Set(servertiming.HeaderName, value)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ServerTimingHeaderValue returns a timing value only for an authenticated admin.
|
||||
func ServerTimingHeaderValue(c *gin.Context) string {
|
||||
if c == nil || c.Request == nil {
|
||||
return ""
|
||||
}
|
||||
role, ok := GetUserRoleFromContext(c)
|
||||
if !ok || role != "admin" {
|
||||
return ""
|
||||
}
|
||||
return servertiming.HeaderValue(c.Request.Context(), time.Now(), responseCacheStatus(c.Writer.Header()))
|
||||
}
|
||||
|
||||
// ServerTimingResponseHeader builds the extra header map required by WebSocket upgrades.
|
||||
func ServerTimingResponseHeader(c *gin.Context) http.Header {
|
||||
value := ServerTimingHeaderValue(c)
|
||||
if value == "" {
|
||||
return nil
|
||||
}
|
||||
return http.Header{servertiming.HeaderName: []string{value}}
|
||||
}
|
||||
|
||||
func isAdminUIRequest(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil || c.Request.URL == nil {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(c.GetHeader(servertiming.AdminUIHeader)) == "1" {
|
||||
return true
|
||||
}
|
||||
path := strings.TrimSpace(c.Request.URL.Path)
|
||||
return path == "/api/v1/admin" || strings.HasPrefix(path, "/api/v1/admin/")
|
||||
}
|
||||
|
||||
func responseCacheStatus(header http.Header) string {
|
||||
for _, name := range []string{snapshotCacheHeader, usageCacheHeader} {
|
||||
switch strings.ToLower(strings.TrimSpace(header.Get(name))) {
|
||||
case "hit":
|
||||
return "hit"
|
||||
case "miss":
|
||||
return "miss"
|
||||
case "bypass":
|
||||
return "bypass"
|
||||
}
|
||||
}
|
||||
return "bypass"
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func runServerTimingRequest(
|
||||
t *testing.T,
|
||||
enabled bool,
|
||||
path string,
|
||||
marker string,
|
||||
role string,
|
||||
handler gin.HandlerFunc,
|
||||
) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := gin.New()
|
||||
engine.Use(ServerTiming(enabled))
|
||||
engine.Any("/*path", func(c *gin.Context) {
|
||||
if role != "" {
|
||||
c.Set(string(ContextKeyUserRole), role)
|
||||
}
|
||||
handler(c)
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
if marker != "" {
|
||||
request.Header.Set(servertiming.AdminUIHeader, marker)
|
||||
}
|
||||
engine.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
|
||||
func TestServerTimingScopesAndRoleGate(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
enabled bool
|
||||
path string
|
||||
marker string
|
||||
role string
|
||||
wantHeader bool
|
||||
}{
|
||||
{name: "disabled", enabled: false, path: "/api/v1/admin/users", role: "admin"},
|
||||
{name: "admin API path", enabled: true, path: "/api/v1/admin/users", role: "admin", wantHeader: true},
|
||||
{name: "shared API marked by admin UI", enabled: true, path: "/api/v1/groups/available", marker: "1", role: "admin", wantHeader: true},
|
||||
{name: "non admin role", enabled: true, path: "/api/v1/groups/available", marker: "1", role: "user"},
|
||||
{name: "unauthenticated public request", enabled: true, path: "/api/v1/settings/public", marker: "1"},
|
||||
{name: "unmarked shared API", enabled: true, path: "/api/v1/groups/available", role: "admin"},
|
||||
{name: "invalid marker", enabled: true, path: "/api/v1/groups/available", marker: "true", role: "admin"},
|
||||
{name: "admin prefix boundary", enabled: true, path: "/api/v1/administrator", role: "admin"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
recorder := runServerTimingRequest(t, tt.enabled, tt.path, tt.marker, tt.role, func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
header := recorder.Header().Get(servertiming.HeaderName)
|
||||
if tt.wantHeader && header == "" {
|
||||
t.Fatalf("%s header missing", servertiming.HeaderName)
|
||||
}
|
||||
if !tt.wantHeader && header != "" {
|
||||
t.Fatalf("unexpected %s header: %q", servertiming.HeaderName, header)
|
||||
}
|
||||
if header != "" && (!strings.Contains(header, "total;dur=") || !strings.Contains(header, `cache;desc="bypass"`)) {
|
||||
t.Fatalf("incomplete timing header: %q", header)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingCollectorIsRequestScoped(t *testing.T) {
|
||||
active := false
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/keys", "1", "admin", func(c *gin.Context) {
|
||||
active = servertiming.Active(c.Request.Context())
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
if !active {
|
||||
t.Fatal("collector was not attached to marked request context")
|
||||
}
|
||||
if recorder.Header().Get(servertiming.HeaderName) == "" {
|
||||
t.Fatal("timing header missing from status-only response")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingFinalizesBeforeEarlyCommit(t *testing.T) {
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/admin/stream", "", "admin", func(c *gin.Context) {
|
||||
c.Status(http.StatusAccepted)
|
||||
c.Writer.WriteHeaderNow()
|
||||
})
|
||||
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
|
||||
t.Fatal("timing header was not written before response commit")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingFinalizesOnFlush(t *testing.T) {
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/admin/export", "", "admin", func(c *gin.Context) {
|
||||
c.Writer.Flush()
|
||||
})
|
||||
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
|
||||
t.Fatal("timing header was not written before stream flush")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingStatusResponses(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
status int
|
||||
}{
|
||||
{name: "not modified", status: http.StatusNotModified},
|
||||
{name: "internal error", status: http.StatusInternalServerError},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/admin/test", "", "admin", func(c *gin.Context) {
|
||||
c.Status(tt.status)
|
||||
})
|
||||
if recorder.Code != tt.status {
|
||||
t.Fatalf("status = %d, want %d", recorder.Code, tt.status)
|
||||
}
|
||||
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
|
||||
t.Fatalf("timing header missing from status %d response", tt.status)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingResponseWriterUnwraps(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
baseWriter := c.Writer
|
||||
writer := &serverTimingResponseWriter{ResponseWriter: baseWriter}
|
||||
if got := writer.Unwrap(); got != baseWriter {
|
||||
t.Fatalf("Unwrap() = %T, want original Gin writer", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingCacheOutcome(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
headerName string
|
||||
value string
|
||||
want string
|
||||
}{
|
||||
{name: "snapshot hit", headerName: snapshotCacheHeader, value: "hit", want: "hit"},
|
||||
{name: "usage miss", headerName: usageCacheHeader, value: "MISS", want: "miss"},
|
||||
{name: "invalid", headerName: snapshotCacheHeader, value: "stale", want: "bypass"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/admin/dashboard", "", "admin", func(c *gin.Context) {
|
||||
c.Header(tt.headerName, tt.value)
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
want := `cache;desc="` + tt.want + `"`
|
||||
if got := recorder.Header().Get(servertiming.HeaderName); !strings.Contains(got, want) {
|
||||
t.Fatalf("timing header %q does not contain %q", got, want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingResponseHeaderForWebSocket(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/ops/ws/qps", nil)
|
||||
collector := servertiming.New(time.Now())
|
||||
c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector))
|
||||
c.Set(string(ContextKeyUserRole), "admin")
|
||||
|
||||
header := ServerTimingResponseHeader(c)
|
||||
if header.Get(servertiming.HeaderName) == "" {
|
||||
t.Fatal("WebSocket response header missing timing value")
|
||||
}
|
||||
|
||||
c.Set(string(ContextKeyUserRole), "user")
|
||||
if got := ServerTimingResponseHeader(c); got != nil {
|
||||
t.Fatalf("non-admin WebSocket received timing header: %#v", got)
|
||||
}
|
||||
}
|
||||
@@ -60,6 +60,7 @@ func SetupRouter(
|
||||
}
|
||||
return nil
|
||||
}))
|
||||
r.Use(middleware2.ServerTiming(cfg.Server.EnableServerTiming))
|
||||
|
||||
// Serve embedded frontend with settings injection if available
|
||||
if web.HasEmbeddedFrontend() {
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
@@ -34,7 +35,7 @@ func newSSRFSafeHTTPClient(timeout time.Duration) *http.Client {
|
||||
TLSHandshakeTimeout: monitorTLSHandshakeTimeout,
|
||||
ResponseHeaderTimeout: monitorResponseHeaderTimeout,
|
||||
}
|
||||
return &http.Client{Timeout: timeout, Transport: tr}
|
||||
return &http.Client{Timeout: timeout, Transport: servertiming.WrapRoundTripper(tr)}
|
||||
}
|
||||
|
||||
// CheckOptions 承载一次检测的自定义入参。
|
||||
|
||||
@@ -22,6 +22,7 @@ import (
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -561,7 +562,7 @@ func NewContentModerationService(
|
||||
userRepo: userRepo,
|
||||
authCacheInvalidator: authCacheInvalidator,
|
||||
emailService: emailService,
|
||||
httpClient: &http.Client{},
|
||||
httpClient: servertiming.InstrumentClient(nil),
|
||||
workerCount: maxContentModerationWorkerCount,
|
||||
asyncQueue: make(chan contentModerationTask, maxContentModerationQueueSize),
|
||||
keyHealth: make(map[string]*contentModerationKeyHealth),
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/payment"
|
||||
"github.com/Wei-Shaw/sub2api/internal/payment/provider"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/shopspring/decimal"
|
||||
)
|
||||
|
||||
@@ -445,7 +446,9 @@ func (s *PaymentService) invokeProvider(ctx context.Context, order *dbent.Paymen
|
||||
IsMobile: req.IsMobile,
|
||||
ReturnURL: providerReturnURL,
|
||||
}, sel, outTradeNo, payAmountStr, subject)
|
||||
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
|
||||
pr, err := prov.CreatePayment(ctx, providerReq)
|
||||
finishProviderCall()
|
||||
if err != nil {
|
||||
slog.Error("[PaymentService] CreatePayment failed", "provider", sel.ProviderKey, "instance", sel.InstanceID, "error", err)
|
||||
if appErr := new(infraerrors.ApplicationError); errors.As(err, &appErr) {
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/ent/paymentorder"
|
||||
"github.com/Wei-Shaw/sub2api/internal/payment"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
)
|
||||
|
||||
// --- Cancel & Expire ---
|
||||
@@ -157,7 +158,9 @@ func (s *PaymentService) checkPaidWithOptions(ctx context.Context, o *dbent.Paym
|
||||
if queryRef == "" {
|
||||
return ""
|
||||
}
|
||||
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
|
||||
resp, err := prov.QueryOrder(ctx, queryRef)
|
||||
finishProviderCall()
|
||||
if err != nil {
|
||||
slog.Warn("query upstream failed", "orderID", o.ID, "error", err)
|
||||
return ""
|
||||
@@ -199,7 +202,9 @@ func (s *PaymentService) checkPaidWithOptions(ctx context.Context, o *dbent.Paym
|
||||
return ""
|
||||
}
|
||||
if cp, ok := prov.(payment.CancelableProvider); ok {
|
||||
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
|
||||
_ = cp.CancelPayment(ctx, queryRef)
|
||||
finishProviderCall()
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -208,7 +213,9 @@ func requeryPaidOrderOnce(ctx context.Context, prov payment.Provider, queryRef s
|
||||
if prov == nil || strings.TrimSpace(queryRef) == "" {
|
||||
return nil, false
|
||||
}
|
||||
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
|
||||
resp, err := prov.QueryOrder(ctx, queryRef)
|
||||
finishProviderCall()
|
||||
if err != nil {
|
||||
slog.Warn("query upstream retry failed", "queryRef", queryRef, "error", err)
|
||||
return nil, false
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/payment"
|
||||
"github.com/Wei-Shaw/sub2api/internal/payment/provider"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
)
|
||||
|
||||
// --- Refund Flow ---
|
||||
@@ -347,12 +348,14 @@ func (s *PaymentService) gwRefund(ctx context.Context, p *RefundPlan) (*payment.
|
||||
})
|
||||
return nil, err
|
||||
}
|
||||
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
|
||||
resp, err := prov.Refund(ctx, payment.RefundRequest{
|
||||
TradeNo: p.Order.PaymentTradeNo,
|
||||
OrderID: p.Order.OutTradeNo,
|
||||
Amount: formatGatewayRefundAmount(p.GatewayAmount, p.Order),
|
||||
Reason: p.Reason,
|
||||
})
|
||||
finishProviderCall()
|
||||
if err != nil {
|
||||
if resp != nil && strings.TrimSpace(resp.Status) == payment.ProviderStatusPending {
|
||||
return resp, nil
|
||||
@@ -417,12 +420,14 @@ func (s *PaymentService) QueryAndFinalizeRefund(ctx context.Context, oid int64)
|
||||
}
|
||||
|
||||
pendingDetail := s.latestRefundPendingDetail(ctx, oid)
|
||||
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
|
||||
resp, err := queryProvider.QueryRefund(ctx, payment.RefundQueryRequest{
|
||||
TradeNo: o.PaymentTradeNo,
|
||||
OrderID: o.OutTradeNo,
|
||||
RefundID: pendingDetail.RefundID,
|
||||
Amount: formatGatewayRefundAmount(o.RefundAmount, o),
|
||||
})
|
||||
finishProviderCall()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query refund: %w", err)
|
||||
}
|
||||
|
||||
@@ -18,6 +18,7 @@ import (
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
@@ -195,7 +196,7 @@ func vertexServiceAccountProxyURL(account *Account) string {
|
||||
func newVertexServiceAccountHTTPClient(proxyURL string) (*http.Client, error) {
|
||||
proxyURL = strings.TrimSpace(proxyURL)
|
||||
if proxyURL == "" {
|
||||
return &http.Client{Timeout: 15 * time.Second}, nil
|
||||
return servertiming.InstrumentClient(&http.Client{Timeout: 15 * time.Second}), nil
|
||||
}
|
||||
|
||||
_, parsedProxy, err := proxyurl.Parse(proxyURL)
|
||||
@@ -211,7 +212,7 @@ func newVertexServiceAccountHTTPClient(proxyURL string) (*http.Client, error) {
|
||||
if err := proxyutil.ConfigureTransportProxy(transport, parsedProxy); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &http.Client{Timeout: 15 * time.Second, Transport: transport}, nil
|
||||
return servertiming.InstrumentClient(&http.Client{Timeout: 15 * time.Second, Transport: transport}), nil
|
||||
}
|
||||
|
||||
func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAccountKey, proxyURL string) (string, time.Duration, error) {
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
@@ -101,6 +102,24 @@ func TestVertexServiceAccountProxyURL(t *testing.T) {
|
||||
require.Empty(t, vertexServiceAccountProxyURL(&Account{ProxyID: &proxyID}))
|
||||
}
|
||||
|
||||
func TestVertexServiceAccountHTTPClientRecordsDependency(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client, err := newVertexServiceAccountHTTPClient("")
|
||||
require.NoError(t, err)
|
||||
collector := servertiming.New(time.Now())
|
||||
ctx := servertiming.WithCollector(context.Background(), collector)
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
|
||||
require.NoError(t, err)
|
||||
response, err := client.Do(request)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, response.Body.Close())
|
||||
require.Contains(t, collector.HeaderValue(time.Now(), "bypass"), "dep_http;dur=")
|
||||
}
|
||||
|
||||
func TestExchangeVertexServiceAccountTokenUsesProxy(t *testing.T) {
|
||||
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -23,6 +23,9 @@ SERVER_PORT=8080
|
||||
# Server mode: release or debug
|
||||
SERVER_MODE=release
|
||||
|
||||
# Return Server-Timing for authenticated requests made by the Admin web UI
|
||||
ENABLE_SERVER_TIMING=false
|
||||
|
||||
# Apple container image overrides (ignored by Docker Compose). Pin release tags
|
||||
# or digests for repeatable operator-managed deployments.
|
||||
APPLE_CONTAINER_SUB2API_IMAGE=weishaw/sub2api:latest
|
||||
|
||||
@@ -20,6 +20,9 @@ server:
|
||||
# Mode: "debug" for development, "release" for production
|
||||
# 运行模式:"debug" 用于开发,"release" 用于生产环境
|
||||
mode: "release"
|
||||
# Return Server-Timing for authenticated requests made by the Admin web UI
|
||||
# 为管理端 Web 页面发出的已认证请求返回 Server-Timing
|
||||
enable_server_timing: false
|
||||
# Frontend base URL used to generate external links in emails (e.g. password reset)
|
||||
# 用于生成邮件中的外部链接(例如:重置密码链接)的前端基础地址
|
||||
# Example: "https://example.com"
|
||||
|
||||
@@ -26,6 +26,7 @@ services:
|
||||
- SERVER_HOST=0.0.0.0
|
||||
- SERVER_PORT=8080
|
||||
- SERVER_MODE=debug
|
||||
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
|
||||
- RUN_MODE=${RUN_MODE:-standard}
|
||||
- DATABASE_HOST=postgres
|
||||
- DATABASE_PORT=5432
|
||||
|
||||
@@ -51,6 +51,7 @@ services:
|
||||
- SERVER_HOST=0.0.0.0
|
||||
- SERVER_PORT=8080
|
||||
- SERVER_MODE=${SERVER_MODE:-release}
|
||||
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
|
||||
- RUN_MODE=${RUN_MODE:-standard}
|
||||
|
||||
# =======================================================================
|
||||
|
||||
@@ -37,6 +37,7 @@ services:
|
||||
- SERVER_HOST=0.0.0.0
|
||||
- SERVER_PORT=8080
|
||||
- SERVER_MODE=${SERVER_MODE:-release}
|
||||
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
|
||||
- RUN_MODE=${RUN_MODE:-standard}
|
||||
|
||||
# =======================================================================
|
||||
|
||||
@@ -47,6 +47,7 @@ services:
|
||||
- SERVER_HOST=0.0.0.0
|
||||
- SERVER_PORT=8080
|
||||
- SERVER_MODE=${SERVER_MODE:-release}
|
||||
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
|
||||
- RUN_MODE=${RUN_MODE:-standard}
|
||||
|
||||
# =======================================================================
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
import { describe, expect, it } from 'vitest'
|
||||
|
||||
import {
|
||||
ADMIN_UI_REQUEST_HEADER,
|
||||
shouldMarkAdminUIRequest,
|
||||
} from '@/api/adminUIRequest'
|
||||
|
||||
describe('Admin UI request marker', () => {
|
||||
it('uses the stable request header name', () => {
|
||||
expect(ADMIN_UI_REQUEST_HEADER).toBe('X-Admin-UI-Request')
|
||||
})
|
||||
|
||||
it.each([
|
||||
'/admin',
|
||||
'/admin/users',
|
||||
'/api/v1/admin',
|
||||
'/api/v1/admin/accounts?status=active',
|
||||
'https://api.example.test/api/v1/admin/dashboard',
|
||||
])('marks Admin API request %s before page navigation', (requestURL) => {
|
||||
expect(shouldMarkAdminUIRequest(requestURL, '/login')).toBe(true)
|
||||
})
|
||||
|
||||
it.each(['/keys', '/groups/available', '/auth/me', '/announcements'])(
|
||||
'marks shared request %s while an Admin page is active',
|
||||
(requestURL) => {
|
||||
expect(shouldMarkAdminUIRequest(requestURL, '/admin/dashboard')).toBe(true)
|
||||
}
|
||||
)
|
||||
|
||||
it.each([
|
||||
['/keys', '/dashboard'],
|
||||
['/api/v1/administer', '/dashboard'],
|
||||
['/keys', '/administrator'],
|
||||
['', '/'],
|
||||
])('does not mark request %s on page %s', (requestURL, pagePath) => {
|
||||
expect(shouldMarkAdminUIRequest(requestURL, pagePath)).toBe(false)
|
||||
})
|
||||
})
|
||||
@@ -12,6 +12,7 @@ describe('API Client', () => {
|
||||
|
||||
beforeEach(async () => {
|
||||
localStorage.clear()
|
||||
window.history.replaceState({}, '', '/')
|
||||
// 每次测试重新导入以获取干净的模块状态
|
||||
vi.resetModules()
|
||||
const mod = await import('@/api/client')
|
||||
@@ -120,6 +121,55 @@ describe('API Client', () => {
|
||||
const config = adapter.mock.calls[0][0]
|
||||
expect(config.withCredentials).toBe(true)
|
||||
})
|
||||
|
||||
it('Admin API 在进入管理页面前也带 Admin UI 标记', async () => {
|
||||
const adapter = vi.fn().mockResolvedValue({
|
||||
status: 200,
|
||||
data: { code: 0, data: {} },
|
||||
headers: {},
|
||||
config: {},
|
||||
statusText: 'OK',
|
||||
})
|
||||
apiClient.defaults.adapter = adapter
|
||||
|
||||
await apiClient.get('/admin/users')
|
||||
|
||||
const config = adapter.mock.calls[0][0]
|
||||
expect(config.headers.get('X-Admin-UI-Request')).toBe('1')
|
||||
})
|
||||
|
||||
it('管理页面调用共享 API 时带 Admin UI 标记', async () => {
|
||||
window.history.replaceState({}, '', '/admin/dashboard')
|
||||
const adapter = vi.fn().mockResolvedValue({
|
||||
status: 200,
|
||||
data: { code: 0, data: {} },
|
||||
headers: {},
|
||||
config: {},
|
||||
statusText: 'OK',
|
||||
})
|
||||
apiClient.defaults.adapter = adapter
|
||||
|
||||
await apiClient.get('/groups/available')
|
||||
|
||||
const config = adapter.mock.calls[0][0]
|
||||
expect(config.headers.get('X-Admin-UI-Request')).toBe('1')
|
||||
})
|
||||
|
||||
it('普通用户页面调用共享 API 时不带 Admin UI 标记', async () => {
|
||||
const adapter = vi.fn().mockResolvedValue({
|
||||
status: 200,
|
||||
data: { code: 0, data: {} },
|
||||
headers: {},
|
||||
config: {},
|
||||
statusText: 'OK',
|
||||
})
|
||||
apiClient.defaults.adapter = adapter
|
||||
|
||||
await apiClient.get('/groups/available')
|
||||
|
||||
const config = adapter.mock.calls[0][0]
|
||||
expect(config.headers.get('X-Admin-UI-Request')).toBeFalsy()
|
||||
})
|
||||
})
|
||||
|
||||
// --- 响应拦截器 ---
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
export const ADMIN_UI_REQUEST_HEADER = 'X-Admin-UI-Request'
|
||||
|
||||
function isAdminPath(path: string): boolean {
|
||||
return (
|
||||
path === '/admin' ||
|
||||
path.startsWith('/admin/') ||
|
||||
path === '/api/v1/admin' ||
|
||||
path.startsWith('/api/v1/admin/')
|
||||
)
|
||||
}
|
||||
|
||||
function requestPath(rawURL: string): string {
|
||||
const value = rawURL.trim()
|
||||
if (!value) return ''
|
||||
try {
|
||||
const origin = typeof window !== 'undefined' ? window.location.origin : 'http://localhost'
|
||||
return new URL(value, origin).pathname
|
||||
} catch {
|
||||
return value.split(/[?#]/, 1)[0]
|
||||
}
|
||||
}
|
||||
|
||||
export function shouldMarkAdminUIRequest(requestURL: string, pagePath?: string): boolean {
|
||||
const currentPath =
|
||||
pagePath ?? (typeof window !== 'undefined' ? window.location.pathname : '')
|
||||
return isAdminPath(requestPath(requestURL)) || isAdminPath(currentPath)
|
||||
}
|
||||
@@ -6,6 +6,7 @@
|
||||
import axios, { AxiosInstance, AxiosError, InternalAxiosRequestConfig, AxiosResponse } from 'axios'
|
||||
import type { ApiResponse } from '@/types'
|
||||
import { getLocale } from '@/i18n'
|
||||
import { ADMIN_UI_REQUEST_HEADER, shouldMarkAdminUIRequest } from './adminUIRequest'
|
||||
import { getAPIBaseURL } from './url'
|
||||
export { buildApiUrl, buildGatewayUrl } from './url'
|
||||
|
||||
@@ -74,6 +75,10 @@ apiClient.interceptors.request.use(
|
||||
config.params.timezone = getUserTimezone()
|
||||
}
|
||||
|
||||
if (config.headers && shouldMarkAdminUIRequest(String(config.url || ''))) {
|
||||
config.headers[ADMIN_UI_REQUEST_HEADER] = '1'
|
||||
}
|
||||
|
||||
return config
|
||||
},
|
||||
(error) => {
|
||||
|
||||
@@ -250,6 +250,7 @@ import TextArea from '@/components/common/TextArea.vue'
|
||||
import { Icon } from '@/components/icons'
|
||||
import { useClipboard } from '@/composables/useClipboard'
|
||||
import { buildApiUrl } from '@/api/client'
|
||||
import { ADMIN_UI_REQUEST_HEADER } from '@/api/adminUIRequest'
|
||||
import { adminAPI } from '@/api/admin'
|
||||
import type { Account, ClaudeModel } from '@/types'
|
||||
|
||||
@@ -438,7 +439,8 @@ const startTest = async () => {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
Authorization: `Bearer ${localStorage.getItem('auth_token')}`,
|
||||
'Content-Type': 'application/json'
|
||||
'Content-Type': 'application/json',
|
||||
[ADMIN_UI_REQUEST_HEADER]: '1'
|
||||
},
|
||||
body: JSON.stringify(requestBody),
|
||||
signal: abortController.signal
|
||||
|
||||
Reference in New Issue
Block a user