package app import ( "bytes" "context" "database/sql" "log" "net/http" "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/platform" "lingniu/vehicle-data-platform/apps/api/internal/static" ) func NewServer(cfg config.Config) http.Handler { var store platform.Store = platform.NewMockStore() if cfg.MySQLDSN != "" { 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) } else { 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") } store = productionStore log.Printf("production mysql store enabled") } } api := platform.NewHandler(platform.NewServiceWithRuntime(store, platform.RuntimeInfo{ RequestTimeoutMs: int(cfg.RequestTimeout / time.Millisecond), })) return withRequestTimeout(static.Handler(cfg.StaticDir, api), cfg.RequestTimeout) } 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()) }