Files
lingniu-vehicle-ingest/vehicle-data-platform/apps/api/internal/app/server.go
2026-07-27 16:46:15 +08:00

657 lines
23 KiB
Go

package app
import (
"bytes"
"container/list"
"context"
"database/sql"
"encoding/json"
"fmt"
"io"
"log"
"math"
"net/http"
"net/url"
"os"
"strconv"
"strings"
"sync"
"time"
"lingniu/vehicle-data-platform/apps/api/internal/config"
"lingniu/vehicle-data-platform/apps/api/internal/httpx"
"lingniu/vehicle-data-platform/apps/api/internal/openplatform"
"lingniu/vehicle-data-platform/apps/api/internal/platform"
"lingniu/vehicle-data-platform/apps/api/internal/static"
)
func NewServer(cfg config.Config) http.Handler {
dataMode := strings.ToLower(strings.TrimSpace(cfg.DataMode))
if dataMode == "" {
dataMode = "mock"
}
var store platform.Store = platform.NewMockStore()
var authDB *sql.DB
var storeErr error
if dataMode != "mock" && dataMode != "production" {
storeErr = fmt.Errorf("DATA_MODE must be mock or production")
}
if dataMode == "production" && strings.TrimSpace(cfg.MySQLDSN) == "" {
storeErr = fmt.Errorf("production data mode requires MYSQL_DSN")
}
if cfg.MySQLDSN != "" && dataMode == "production" {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
db, err := platform.OpenSQL(ctx, "mysql", cfg.MySQLDSN)
if err != nil {
log.Printf("production mysql store disabled: %v", err)
if dataMode == "production" {
storeErr = fmt.Errorf("connect production mysql: %w", err)
}
} else {
authDB = db
var tdengine *sql.DB
if cfg.TDengineDSN != "" {
tdengine, err = platform.OpenSQL(ctx, cfg.TDengineDriver, cfg.TDengineDSN)
if err != nil {
log.Printf("production tdengine store disabled: %v", err)
} else {
log.Printf("production tdengine store enabled")
}
}
productionStore := platform.NewProductionStore(db, tdengine, cfg.TDengineDatabase)
if cfg.RedisAddr != "" {
productionStore.WithRedisOnlineKeyCounter(platform.NewRedisOnlineKeyCounter(cfg.RedisAddr, cfg.RedisUsername, cfg.RedisPassword, cfg.RedisDB))
log.Printf("production redis health probe enabled")
}
if cfg.CapacityCheckBin != "" {
productionStore.WithCapacityChecker(platform.NewCapacityCheckCommand(cfg.CapacityCheckBin))
log.Printf("production capacity-check probe enabled")
}
productionStore.WithAlertStreamConfig(cfg.AlertStreamMode, cfg.AlertStreamKafkaGroup)
store = productionStore
storeErr = nil
log.Printf("production mysql store enabled")
}
}
host, _ := os.Hostname()
workerHost := strings.TrimSpace(host)
if workerHost == "" {
workerHost = "api"
}
platformService := platform.NewServiceWithRuntime(store, platform.RuntimeInfo{
DataMode: dataMode,
ExportDir: strings.TrimSpace(cfg.ExportDir),
RequestTimeoutMs: int(cfg.RequestTimeout / time.Millisecond),
AMapWebJSConfigured: strings.TrimSpace(cfg.AMapWebJSKey) != "",
AMapAPIConfigured: strings.TrimSpace(cfg.AMapAPIKey) != "",
AMapSecurityProxyEnabled: strings.TrimSpace(cfg.AMapSecurityCode) != "" && strings.TrimSpace(cfg.AMapServiceHost) != "",
AMapSecurityCodeExposed: exposedAMapSecurityCode(cfg) != "",
AMapSecurityServiceHost: strings.TrimSpace(cfg.AMapServiceHost),
PlatformRelease: strings.TrimSpace(cfg.PlatformRelease),
AlertStreamMode: strings.TrimSpace(cfg.AlertStreamMode),
AlertStreamConsumerGroup: strings.TrimSpace(cfg.AlertStreamKafkaGroup),
AlertNotificationConfig: alertNotificationConfig(cfg, dataMode),
HistoryCleanupAutomation: cfg.HistoryCleanupAutomation,
HistoryCleanupPoll: cfg.HistoryCleanupPollInterval,
HistoryCleanupLease: cfg.HistoryCleanupLease,
HistoryCleanupWorkerID: "cleanup-scheduler-" + workerHost,
})
platformService.StartHistoryExportCleanupAutomation(context.Background())
var api http.Handler = platform.NewHandler(platformService)
if storeErr != nil {
log.Printf("platform data store unavailable: %v", storeErr)
api = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
httpx.WriteError(w, http.StatusServiceUnavailable, "DATA_STORE_UNAVAILABLE", "生产数据源不可用", storeErr.Error(), requestTraceID(r))
})
}
if authDB != nil && storeErr == nil {
openPlatformHandler := openplatform.NewAdminHandler(
openplatform.NewService(openplatform.NewMySQLRepository(authDB)),
openplatform.NewPortalService(authDB, cfg.SessionTTL),
)
internalMux := http.NewServeMux()
internalMux.Handle("/api/v2/open-platform/", openPlatformHandler)
internalMux.Handle("/", api)
api = internalMux
}
// Reverse geocoding consumes the server-side AMap credential, so it must
// stay behind the same authentication boundary as the platform API.
api = withAMapReverseGeocodeAPI(api, cfg, "https://restapi.amap.com", http.DefaultClient)
handler := static.Handler(cfg.StaticDir, withAPIAuth(api, cfg, authDB))
handler = withAppConfig(handler, cfg)
handler = withAMapSecurityProxy(handler, cfg, defaultAMapProxyUpstreams(), http.DefaultClient)
return withRequestTimeout(handler, cfg.RequestTimeout)
}
func alertNotificationConfig(cfg config.Config, dataMode string) platform.AlertNotificationConfig {
channels := []platform.AlertNotificationChannelCapability{
{Channel: "in_app", Label: "站内信", Configured: true},
{Channel: "sms", Label: "短信", Configured: dataMode == "mock" || strings.TrimSpace(cfg.AlertNotificationSMSURL) != "" && strings.TrimSpace(cfg.AlertNotificationSMSSecret) != ""},
{Channel: "email", Label: "邮件", Configured: dataMode == "mock" || strings.TrimSpace(cfg.AlertNotificationEmailURL) != "" && strings.TrimSpace(cfg.AlertNotificationEmailSecret) != ""},
{Channel: "wecom", Label: "企业通讯", Configured: dataMode == "mock" || strings.TrimSpace(cfg.AlertNotificationWeComURL) != "" && strings.TrimSpace(cfg.AlertNotificationWeComSecret) != ""},
}
targets := []platform.AlertNotificationTargetOption{{ID: "platform-operators", Label: "平台值班组", Channels: []string{"in_app"}}}
raw := strings.TrimSpace(cfg.AlertNotificationTargetsJSON)
if raw == "" && dataMode == "mock" {
targets = []platform.AlertNotificationTargetOption{
{ID: "platform-operators", Label: "平台值班组", Channels: []string{"in_app", "email", "wecom"}},
{ID: "night-shift", Label: "夜班负责人", Channels: []string{"in_app", "sms", "wecom"}},
{ID: "data-platform", Label: "数据平台负责人", Channels: []string{"in_app", "email", "wecom"}},
}
} else if raw != "" {
var configured []platform.AlertNotificationTargetOption
if err := json.Unmarshal([]byte(raw), &configured); err != nil {
log.Printf("alert notification target catalog ignored: %v", err)
} else {
targets = configured
}
}
return platform.NormalizeAlertNotificationConfig(platform.AlertNotificationConfig{Targets: targets, Channels: channels})
}
func withAppConfig(next http.Handler, cfg config.Config) http.Handler {
type appConfig struct {
AMapWebJSKey string `json:"amapWebJsKey,omitempty"`
AMapSecurityCode string `json:"amapSecurityJsCode,omitempty"`
AMapServiceHost string `json:"amapSecurityServiceHost,omitempty"`
}
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/app-config.js" {
next.ServeHTTP(w, r)
return
}
w.Header().Set("Content-Type", "application/javascript; charset=utf-8")
w.Header().Set("Cache-Control", "no-store")
body, err := json.Marshal(appConfig{
AMapWebJSKey: cfg.AMapWebJSKey,
AMapSecurityCode: exposedAMapSecurityCode(cfg),
AMapServiceHost: cfg.AMapServiceHost,
})
if err != nil {
http.Error(w, "failed to render app config", http.StatusInternalServerError)
return
}
_, _ = fmt.Fprintf(w, "window.__LINGNIU_APP_CONFIG__=%s;\n", body)
})
}
func exposedAMapSecurityCode(cfg config.Config) string {
if cfg.AMapServiceHost != "" {
return ""
}
return cfg.AMapSecurityCode
}
type amapProxyUpstreams struct {
RestAPI string
WebAPI string
FMap string
}
func defaultAMapProxyUpstreams() amapProxyUpstreams {
return amapProxyUpstreams{
RestAPI: "https://restapi.amap.com",
WebAPI: "https://webapi.amap.com",
FMap: "https://fmap01.amap.com",
}
}
func (upstreams amapProxyUpstreams) normalized() amapProxyUpstreams {
defaults := defaultAMapProxyUpstreams()
if upstreams.RestAPI == "" {
upstreams.RestAPI = defaults.RestAPI
}
if upstreams.WebAPI == "" {
upstreams.WebAPI = defaults.WebAPI
}
if upstreams.FMap == "" {
upstreams.FMap = defaults.FMap
}
upstreams.RestAPI = strings.TrimRight(upstreams.RestAPI, "/")
upstreams.WebAPI = strings.TrimRight(upstreams.WebAPI, "/")
upstreams.FMap = strings.TrimRight(upstreams.FMap, "/")
return upstreams
}
func withAMapSecurityProxy(next http.Handler, cfg config.Config, upstreams amapProxyUpstreams, client *http.Client) http.Handler {
serviceHost := strings.TrimRight(cfg.AMapServiceHost, "/")
securityCode := strings.TrimSpace(cfg.AMapSecurityCode)
if serviceHost == "" || securityCode == "" {
return next
}
if client == nil {
client = http.DefaultClient
}
upstreams = upstreams.normalized()
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != serviceHost && !strings.HasPrefix(r.URL.Path, serviceHost+"/") {
next.ServeHTTP(w, r)
return
}
target, err := amapProxyURL(upstreams, serviceHost, securityCode, r)
if err != nil {
httpx.WriteError(w, http.StatusBadGateway, "AMAP_PROXY_URL_INVALID", "高德地图代理地址无效", err.Error(), requestTraceID(r))
return
}
req, err := http.NewRequestWithContext(r.Context(), r.Method, target, r.Body)
if err != nil {
httpx.WriteError(w, http.StatusBadGateway, "AMAP_PROXY_REQUEST_INVALID", "高德地图代理请求无效", err.Error(), requestTraceID(r))
return
}
req.Header = r.Header.Clone()
req.Host = ""
resp, err := client.Do(req)
if err != nil {
httpx.WriteError(w, http.StatusBadGateway, "AMAP_PROXY_FAILED", "高德地图代理请求失败", err.Error(), requestTraceID(r))
return
}
defer resp.Body.Close()
for key, values := range resp.Header {
for _, value := range values {
w.Header().Add(key, value)
}
}
w.WriteHeader(resp.StatusCode)
_, _ = io.Copy(w, resp.Body)
})
}
func amapProxyURL(upstreams amapProxyUpstreams, serviceHost string, securityCode string, r *http.Request) (string, error) {
targetPath := strings.TrimPrefix(r.URL.Path, serviceHost)
if targetPath == "" {
targetPath = "/"
}
upstreamBase := upstreams.RestAPI
if strings.HasPrefix(targetPath, "/v4/map/styles") {
upstreamBase = upstreams.WebAPI
} else if strings.HasPrefix(targetPath, "/v3/vectormap") {
upstreamBase = upstreams.FMap
}
base, err := url.Parse(upstreamBase)
if err != nil {
return "", err
}
base.Path = strings.TrimRight(base.Path, "/") + targetPath
query := r.URL.Query()
query.Set("jscode", securityCode)
base.RawQuery = query.Encode()
return base.String(), nil
}
type mapReverseGeocodeResponse struct {
Provider string `json:"provider"`
Longitude float64 `json:"longitude"`
Latitude float64 `json:"latitude"`
FormattedAddress string `json:"formattedAddress"`
Province string `json:"province,omitempty"`
City string `json:"city,omitempty"`
District string `json:"district,omitempty"`
Township string `json:"township,omitempty"`
Adcode string `json:"adcode,omitempty"`
}
const (
reverseGeocodeCacheTTL = time.Hour
reverseGeocodeCacheMaxEntries = 4096
)
type reverseGeocodeCacheEntry struct {
key string
value mapReverseGeocodeResponse
expiresAt time.Time
}
type reverseGeocodeKeyLock struct {
mu sync.Mutex
refs int
}
type reverseGeocodeCache struct {
mu sync.Mutex
entries map[string]*list.Element
recency *list.List
keyLocks map[string]*reverseGeocodeKeyLock
maxEntries int
ttl time.Duration
now func() time.Time
}
func newReverseGeocodeCache(maxEntries int, ttl time.Duration) *reverseGeocodeCache {
return &reverseGeocodeCache{
entries: make(map[string]*list.Element),
recency: list.New(),
keyLocks: make(map[string]*reverseGeocodeKeyLock),
maxEntries: maxEntries,
ttl: ttl,
now: time.Now,
}
}
func reverseGeocodeCacheKey(longitude, latitude float64) string {
return fmt.Sprintf("%.5f,%.5f", longitude, latitude)
}
func (cache *reverseGeocodeCache) get(key string) (mapReverseGeocodeResponse, bool) {
cache.mu.Lock()
defer cache.mu.Unlock()
element, ok := cache.entries[key]
if !ok {
return mapReverseGeocodeResponse{}, false
}
entry := element.Value.(*reverseGeocodeCacheEntry)
if !entry.expiresAt.After(cache.now()) {
cache.remove(element)
return mapReverseGeocodeResponse{}, false
}
cache.recency.MoveToFront(element)
return entry.value, true
}
func (cache *reverseGeocodeCache) put(key string, value mapReverseGeocodeResponse) {
if cache.maxEntries <= 0 || cache.ttl <= 0 {
return
}
cache.mu.Lock()
defer cache.mu.Unlock()
if element, ok := cache.entries[key]; ok {
entry := element.Value.(*reverseGeocodeCacheEntry)
entry.value = value
entry.expiresAt = cache.now().Add(cache.ttl)
cache.recency.MoveToFront(element)
return
}
for len(cache.entries) >= cache.maxEntries {
cache.remove(cache.recency.Back())
}
entry := &reverseGeocodeCacheEntry{key: key, value: value, expiresAt: cache.now().Add(cache.ttl)}
cache.entries[key] = cache.recency.PushFront(entry)
}
func (cache *reverseGeocodeCache) remove(element *list.Element) {
if element == nil {
return
}
delete(cache.entries, element.Value.(*reverseGeocodeCacheEntry).key)
cache.recency.Remove(element)
}
func (cache *reverseGeocodeCache) lockKey(key string) func() {
cache.mu.Lock()
keyLock := cache.keyLocks[key]
if keyLock == nil {
keyLock = &reverseGeocodeKeyLock{}
cache.keyLocks[key] = keyLock
}
keyLock.refs++
cache.mu.Unlock()
keyLock.mu.Lock()
return func() {
keyLock.mu.Unlock()
cache.mu.Lock()
keyLock.refs--
if keyLock.refs == 0 {
delete(cache.keyLocks, key)
}
cache.mu.Unlock()
}
}
func withAMapReverseGeocodeAPI(next http.Handler, cfg config.Config, upstream string, client *http.Client) http.Handler {
apiKey := strings.TrimSpace(cfg.AMapAPIKey)
if client == nil {
client = http.DefaultClient
}
upstream = strings.TrimRight(upstream, "/")
if upstream == "" {
upstream = "https://restapi.amap.com"
}
cache := newReverseGeocodeCache(reverseGeocodeCacheMaxEntries, reverseGeocodeCacheTTL)
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/map/reverse-geocode" {
next.ServeHTTP(w, r)
return
}
if r.Method != http.MethodGet {
httpx.WriteError(w, http.StatusMethodNotAllowed, "METHOD_NOT_ALLOWED", "请求方法不支持", "", requestTraceID(r))
return
}
if apiKey == "" {
httpx.WriteError(w, http.StatusServiceUnavailable, "AMAP_API_UNCONFIGURED", "高德服务端 API 未配置", "请配置 AMAP_API_KEY 后再解析地址", requestTraceID(r))
return
}
longitude, latitude, err := parseReverseGeocodeCoordinate(r.URL.Query())
if err != nil {
httpx.WriteError(w, http.StatusBadRequest, "BAD_COORDINATE", "经纬度参数无效", err.Error(), requestTraceID(r))
return
}
cacheKey := reverseGeocodeCacheKey(longitude, latitude)
unlock := cache.lockKey(cacheKey)
defer unlock()
if cached, ok := cache.get(cacheKey); ok {
cached.Longitude = longitude
cached.Latitude = latitude
w.Header().Set("X-Reverse-Geocode-Cache", "HIT")
httpx.WriteOK(w, requestTraceID(r), cached)
return
}
w.Header().Set("X-Reverse-Geocode-Cache", "MISS")
target, err := url.Parse(upstream + "/v3/geocode/regeo")
if err != nil {
httpx.WriteError(w, http.StatusBadGateway, "AMAP_REVERSE_GEOCODE_URL_INVALID", "高德逆地理编码地址无效", err.Error(), requestTraceID(r))
return
}
query := target.Query()
query.Set("key", apiKey)
mapLongitude, mapLatitude := wgs84ToGCJ02(longitude, latitude)
query.Set("location", fmt.Sprintf("%.6f,%.6f", mapLongitude, mapLatitude))
query.Set("extensions", "base")
query.Set("radius", "1000")
query.Set("output", "JSON")
target.RawQuery = query.Encode()
req, err := http.NewRequestWithContext(r.Context(), http.MethodGet, target.String(), nil)
if err != nil {
httpx.WriteError(w, http.StatusBadGateway, "AMAP_REVERSE_GEOCODE_REQUEST_INVALID", "高德逆地理编码请求无效", err.Error(), requestTraceID(r))
return
}
resp, err := client.Do(req)
if err != nil {
httpx.WriteError(w, http.StatusBadGateway, "AMAP_REVERSE_GEOCODE_FAILED", "高德逆地理编码请求失败", err.Error(), requestTraceID(r))
return
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
httpx.WriteError(w, http.StatusBadGateway, "AMAP_REVERSE_GEOCODE_BAD_STATUS", "高德逆地理编码返回异常", resp.Status, requestTraceID(r))
return
}
var body struct {
Status string `json:"status"`
Info string `json:"info"`
Infocode string `json:"infocode"`
Regeocode struct {
FormattedAddress string `json:"formatted_address"`
AddressComponent struct {
Province json.RawMessage `json:"province"`
City json.RawMessage `json:"city"`
District json.RawMessage `json:"district"`
Township json.RawMessage `json:"township"`
Adcode json.RawMessage `json:"adcode"`
} `json:"addressComponent"`
} `json:"regeocode"`
}
if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
httpx.WriteError(w, http.StatusBadGateway, "AMAP_REVERSE_GEOCODE_BAD_JSON", "高德逆地理编码响应解析失败", err.Error(), requestTraceID(r))
return
}
if body.Status != "1" {
detail := strings.TrimSpace(body.Info)
if body.Infocode != "" {
detail = strings.TrimSpace(detail + " " + body.Infocode)
}
httpx.WriteError(w, http.StatusBadGateway, "AMAP_REVERSE_GEOCODE_REJECTED", "高德逆地理编码未成功", detail, requestTraceID(r))
return
}
result := mapReverseGeocodeResponse{
Provider: "AMap",
Longitude: longitude,
Latitude: latitude,
FormattedAddress: body.Regeocode.FormattedAddress,
Province: amapJSONText(body.Regeocode.AddressComponent.Province),
City: amapJSONText(body.Regeocode.AddressComponent.City),
District: amapJSONText(body.Regeocode.AddressComponent.District),
Township: amapJSONText(body.Regeocode.AddressComponent.Township),
Adcode: amapJSONText(body.Regeocode.AddressComponent.Adcode),
}
cache.put(cacheKey, result)
httpx.WriteOK(w, requestTraceID(r), result)
})
}
func parseReverseGeocodeCoordinate(query url.Values) (float64, float64, error) {
longitude, err := strconv.ParseFloat(strings.TrimSpace(firstNonEmpty(query.Get("longitude"), query.Get("lng"))), 64)
if err != nil {
return 0, 0, fmt.Errorf("longitude 必须是数字")
}
latitude, err := strconv.ParseFloat(strings.TrimSpace(firstNonEmpty(query.Get("latitude"), query.Get("lat"))), 64)
if err != nil {
return 0, 0, fmt.Errorf("latitude 必须是数字")
}
if !isCoordinate(longitude, -180, 180) || !isCoordinate(latitude, -90, 90) || (longitude == 0 && latitude == 0) {
return 0, 0, fmt.Errorf("longitude/latitude 超出范围或为空坐标")
}
return longitude, latitude, nil
}
func wgs84ToGCJ02(longitude float64, latitude float64) (float64, float64) {
if longitude < 72.004 || longitude > 137.8347 || latitude < 0.8293 || latitude > 55.8271 {
return longitude, latitude
}
const semiMajorAxis = 6378245.0
const eccentricitySquared = 0.006693421622965943
longitudeOffset := transformGCJLongitude(longitude-105, latitude-35)
latitudeOffset := transformGCJLatitude(longitude-105, latitude-35)
radianLatitude := latitude / 180 * math.Pi
magic := 1 - eccentricitySquared*math.Pow(math.Sin(radianLatitude), 2)
squareRootMagic := math.Sqrt(magic)
convertedLatitude := latitude + latitudeOffset*180/((semiMajorAxis*(1-eccentricitySquared))/(magic*squareRootMagic)*math.Pi)
convertedLongitude := longitude + longitudeOffset*180/(semiMajorAxis/squareRootMagic*math.Cos(radianLatitude)*math.Pi)
return convertedLongitude, convertedLatitude
}
func transformGCJLatitude(longitude float64, latitude float64) float64 {
value := -100 + 2*longitude + 3*latitude + 0.2*latitude*latitude + 0.1*longitude*latitude + 0.2*math.Sqrt(math.Abs(longitude))
value += (20*math.Sin(6*longitude*math.Pi) + 20*math.Sin(2*longitude*math.Pi)) * 2 / 3
value += (20*math.Sin(latitude*math.Pi) + 40*math.Sin(latitude/3*math.Pi)) * 2 / 3
value += (160*math.Sin(latitude/12*math.Pi) + 320*math.Sin(latitude*math.Pi/30)) * 2 / 3
return value
}
func transformGCJLongitude(longitude float64, latitude float64) float64 {
value := 300 + longitude + 2*latitude + 0.1*longitude*longitude + 0.1*longitude*latitude + 0.1*math.Sqrt(math.Abs(longitude))
value += (20*math.Sin(6*longitude*math.Pi) + 20*math.Sin(2*longitude*math.Pi)) * 2 / 3
value += (20*math.Sin(longitude*math.Pi) + 40*math.Sin(longitude/3*math.Pi)) * 2 / 3
value += (150*math.Sin(longitude/12*math.Pi) + 300*math.Sin(longitude/30*math.Pi)) * 2 / 3
return value
}
func isCoordinate(value float64, min float64, max float64) bool {
return !math.IsNaN(value) && !math.IsInf(value, 0) && value >= min && value <= max
}
func amapJSONText(value json.RawMessage) string {
if len(value) == 0 || string(value) == "null" {
return ""
}
var text string
if err := json.Unmarshal(value, &text); err == nil {
return text
}
var texts []string
if err := json.Unmarshal(value, &texts); err == nil {
return strings.Join(texts, "")
}
return ""
}
func firstNonEmpty(values ...string) string {
for _, value := range values {
if strings.TrimSpace(value) != "" {
return value
}
}
return ""
}
func withRequestTimeout(next http.Handler, timeout time.Duration) http.Handler {
if timeout <= 0 {
return next
}
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ctx, cancel := context.WithTimeout(r.Context(), timeout)
defer cancel()
buffered := newBufferedResponseWriter()
result := make(chan any, 1)
go func() {
defer func() {
result <- recover()
}()
next.ServeHTTP(buffered, r.WithContext(ctx))
}()
select {
case panicValue := <-result:
if panicValue != nil {
panic(panicValue)
}
buffered.writeTo(w)
case <-ctx.Done():
httpx.WriteError(w, http.StatusServiceUnavailable, "REQUEST_TIMEOUT", "请求处理超时", "", requestTraceID(r))
}
})
}
func requestTraceID(r *http.Request) string {
if value := r.Header.Get("X-Trace-Id"); value != "" {
return value
}
return "trace-" + time.Now().Format("20060102150405.000000")
}
type bufferedResponseWriter struct {
mu sync.Mutex
header http.Header
status int
body bytes.Buffer
}
func newBufferedResponseWriter() *bufferedResponseWriter {
return &bufferedResponseWriter{header: http.Header{}, status: http.StatusOK}
}
func (w *bufferedResponseWriter) Header() http.Header {
return w.header
}
func (w *bufferedResponseWriter) WriteHeader(statusCode int) {
w.mu.Lock()
defer w.mu.Unlock()
w.status = statusCode
}
func (w *bufferedResponseWriter) Write(body []byte) (int, error) {
w.mu.Lock()
defer w.mu.Unlock()
return w.body.Write(body)
}
func (w *bufferedResponseWriter) writeTo(target http.ResponseWriter) {
w.mu.Lock()
defer w.mu.Unlock()
for key, values := range w.header {
for _, value := range values {
target.Header().Add(key, value)
}
}
target.WriteHeader(w.status)
_, _ = target.Write(w.body.Bytes())
}