lingniu-vehicle-ingest/vehicle-data-platform/apps/api/internal/app/server.go

package app

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"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(withAppConfig(static.Handler(cfg.StaticDir, api), cfg), cfg.RequestTimeout)
}

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: cfg.AMapSecurityCode,
			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 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())
}