AIRouter/internal/server/idempotency.go

package server

import (
	"net/http"
	"time"

	"github.com/edgeai/gateway/internal/handler"
)

// idempotencyEntry 幂等键缓存条目。
// 记录请求处理状态和响应，用于在重复请求时返回缓存结果。
type idempotencyEntry struct {
	status    int
	response  []byte
	createdAt time.Time
	inFlight  bool // 正在处理中
}

const (
	// idempotencyTTL 幂等键缓存存活时间。
	idempotencyTTL = 10 * time.Minute
	// idempotencyCleanupInterval 清理间隔。
	idempotencyCleanupInterval = 5 * time.Minute
)

// checkIdempotency 检查幂等键，如果重复请求则返回缓存的响应。
// 如果是首次请求，返回 nil 表示可以继续处理。
// 如果是正在处理中的重复请求，返回 409 Conflict。
func (s *Server) checkIdempotency(w http.ResponseWriter, key, appID string) bool {
	if key == "" {
		return false
	}

	cacheKey := appID + ":" + key

	s.idempotencyMu.Lock()
	entry, exists := s.idempotencyCache[cacheKey]
	if exists {
		if entry.inFlight {
			// 正在处理中，返回 409
			s.idempotencyMu.Unlock()
			handler.WriteError(w, handler.NewGatewayError(handler.ErrInvalidRequest, "duplicate idempotency key, request already in progress"))
			return true
		}

		// 检查是否过期
		if time.Since(entry.createdAt) > idempotencyTTL {
			delete(s.idempotencyCache, cacheKey)
		} else {
			// 返回缓存的响应
			s.idempotencyMu.Unlock()
			w.Header().Set("Content-Type", "application/json")
			w.Header().Set("X-Idempotent-Replay", "true")
			w.WriteHeader(entry.status)
			w.Write(entry.response)
			return true
		}
	}

	// 标记为正在处理中
	s.idempotencyCache[cacheKey] = &idempotencyEntry{
		inFlight:  true,
		createdAt: time.Now(),
	}
	s.idempotencyMu.Unlock()
	return false
}

// storeIdempotencyResult 存储幂等请求的响应结果。
func (s *Server) storeIdempotencyResult(key, appID string, status int, response []byte) {
	if key == "" {
		return
	}

	cacheKey := appID + ":" + key

	s.idempotencyMu.Lock()
	s.idempotencyCache[cacheKey] = &idempotencyEntry{
		status:    status,
		response:  response,
		createdAt: time.Now(),
		inFlight:  false,
	}
	s.idempotencyMu.Unlock()
}

// cleanupIdempotencyCache 清理过期的幂等键缓存。
func (s *Server) cleanupIdempotencyCache() {
	s.idempotencyMu.Lock()
	defer s.idempotencyMu.Unlock()

	now := time.Now()
	for k, entry := range s.idempotencyCache {
		if now.Sub(entry.createdAt) > idempotencyTTL {
			delete(s.idempotencyCache, k)
		}
	}
}

// startIdempotencyCleanup 启动后台清理任务，定期清理过期的幂等键。
// 返回停止函数。
func (s *Server) startIdempotencyCleanup() func() {
	ticker := time.NewTicker(idempotencyCleanupInterval)
	stopCh := make(chan struct{})

	go func() {
		for {
			select {
			case <-ticker.C:
				s.cleanupIdempotencyCache()
			case <-stopCh:
				ticker.Stop()
				return
			}
		}
	}()

	return func() {
		close(stopCh)
	}
}

// captureResponseWriter 捕获响应状态和响应体，用于幂等性缓存。
type captureResponseWriter struct {
	http.ResponseWriter
	statusCode int
	body       []byte
}

func (crw *captureResponseWriter) WriteHeader(code int) {
	crw.statusCode = code
	crw.ResponseWriter.WriteHeader(code)
}

func (crw *captureResponseWriter) Write(b []byte) (int, error) {
	if crw.statusCode == 0 {
		crw.statusCode = 200
	}
	crw.body = append(crw.body, b...)
	return crw.ResponseWriter.Write(b)
}

func (crw *captureResponseWriter) Flush() {
	if f, ok := crw.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}