feat(relay): complete Responses WebSocket v2 support

common/json.go

@@ -6,6 +6,8 @@ import (
 	"io"
 )
 
+type RawMessage = json.RawMessage
+
 func Unmarshal(data []byte, v any) error {
 	return json.Unmarshal(data, v)
 }
@@ -22,7 +24,7 @@ func Marshal(v any) ([]byte, error) {
 	return json.Marshal(v)
 }
 
-func GetJsonType(data json.RawMessage) string {
+func GetJsonType(data RawMessage) string {
 	trimmed := bytes.TrimSpace(data)
 	if len(trimmed) == 0 {
 		return "unknown"
@@ -45,7 +47,7 @@ func GetJsonType(data json.RawMessage) string {
 }
 
 // JsonRawMessageToString returns JSON strings as their decoded value and other JSON values as raw text.
-func JsonRawMessageToString(data json.RawMessage) string {
+func JsonRawMessageToString(data RawMessage) string {
 	trimmed := bytes.TrimSpace(data)
 	if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) {
 		return ""

common/rate-limit.go

@@ -6,22 +6,26 @@ import (
 )
 
 type InMemoryRateLimiter struct {
-	store              map[string]*[]int64
+	store              map[string]*[]rateLimitEntry
 	mutex              sync.Mutex
 	expirationDuration time.Duration
 }
 
+type rateLimitEntry struct {
+	timestamp     int64
+	reservationID string
+}
+
 func (l *InMemoryRateLimiter) Init(expirationDuration time.Duration) {
-	if l.store == nil {
-		l.mutex.Lock()
-		if l.store == nil {
-			l.store = make(map[string]*[]int64)
-			l.expirationDuration = expirationDuration
-			if expirationDuration > 0 {
-				go l.clearExpiredItems()
-			}
-		}
-		l.mutex.Unlock()
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+	if l.store != nil {
+		return
+	}
+	l.store = make(map[string]*[]rateLimitEntry)
+	l.expirationDuration = expirationDuration
+	if expirationDuration > 0 {
+		go l.clearExpiredItems()
 	}
 }
 
@@ -33,7 +37,7 @@ func (l *InMemoryRateLimiter) clearExpiredItems() {
 		for key := range l.store {
 			queue := l.store[key]
 			size := len(*queue)
-			if size == 0 || now-(*queue)[size-1] > int64(l.expirationDuration.Seconds()) {
+			if size == 0 || now-(*queue)[size-1].timestamp > int64(l.expirationDuration.Seconds()) {
 				delete(l.store, key)
 			}
 		}
@@ -45,26 +49,96 @@ func (l *InMemoryRateLimiter) clearExpiredItems() {
 func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration int64) bool {
 	l.mutex.Lock()
 	defer l.mutex.Unlock()
+	return l.requestLocked(key, maxRequestNum, duration, "")
+}
+
+// Reserve atomically records a provisional request. The caller must later
+// commit or release the reservation.
+func (l *InMemoryRateLimiter) Reserve(key string, maxRequestNum int, duration int64, reservationID string) bool {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+	return l.requestLocked(key, maxRequestNum, duration, reservationID)
+}
+
+func (l *InMemoryRateLimiter) requestLocked(key string, maxRequestNum int, duration int64, reservationID string) bool {
+	if maxRequestNum == 0 {
+		return true
+	}
 	// [old <-- new]
 	queue, ok := l.store[key]
 	now := time.Now().Unix()
+	entry := rateLimitEntry{timestamp: now, reservationID: reservationID}
 	if ok {
 		if len(*queue) < maxRequestNum {
-			*queue = append(*queue, now)
+			*queue = append(*queue, entry)
 			return true
 		} else {
-			if now-(*queue)[0] >= duration {
+			if now-(*queue)[0].timestamp >= duration {
 				*queue = (*queue)[1:]
-				*queue = append(*queue, now)
+				*queue = append(*queue, entry)
 				return true
 			} else {
 				return false
 			}
 		}
 	} else {
-		s := make([]int64, 0, maxRequestNum)
+		s := make([]rateLimitEntry, 0, maxRequestNum)
 		l.store[key] = &s
-		*(l.store[key]) = append(*(l.store[key]), now)
+		*(l.store[key]) = append(*(l.store[key]), entry)
 	}
 	return true
 }
+
+// CommitReservation finalizes a provisional request while preserving the
+// admission timestamp used by the sliding window.
+func (l *InMemoryRateLimiter) CommitReservation(key string, reservationID string) bool {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+	queue, ok := l.store[key]
+	if !ok {
+		return false
+	}
+	for i := range *queue {
+		if (*queue)[i].reservationID == reservationID {
+			(*queue)[i].reservationID = ""
+			return true
+		}
+	}
+	return false
+}
+
+// ReleaseReservation removes a provisional request that did not succeed.
+func (l *InMemoryRateLimiter) ReleaseReservation(key string, reservationID string) bool {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+	queue, ok := l.store[key]
+	if !ok {
+		return false
+	}
+	for i := range *queue {
+		if (*queue)[i].reservationID == reservationID {
+			*queue = append((*queue)[:i], (*queue)[i+1:]...)
+			if len(*queue) == 0 {
+				delete(l.store, key)
+			}
+			return true
+		}
+	}
+	return false
+}
+
+// Check reports whether a request would be allowed without recording it.
+// The duration parameter's unit is seconds.
+func (l *InMemoryRateLimiter) Check(key string, maxRequestNum int, duration int64) bool {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+	if maxRequestNum == 0 {
+		return true
+	}
+	queue, ok := l.store[key]
+	if !ok || len(*queue) < maxRequestNum {
+		return true
+	}
+	now := time.Now().Unix()
+	return now-(*queue)[0].timestamp >= duration
+}

common/rate_limit_test.go

@@ -0,0 +1,31 @@
+package common
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInMemoryRateLimiterConcurrentInit(t *testing.T) {
+	var limiter InMemoryRateLimiter
+	start := make(chan struct{})
+	var ready sync.WaitGroup
+	var done sync.WaitGroup
+	ready.Add(2)
+	done.Add(2)
+	for range 2 {
+		go func() {
+			defer done.Done()
+			ready.Done()
+			<-start
+			limiter.Init(0)
+		}()
+	}
+	ready.Wait()
+	close(start)
+	done.Wait()
+
+	assert.True(t, limiter.Request("test", 1, 60))
+	assert.False(t, limiter.Request("test", 1, 60))
+}

controller/channel.go

@@ -71,6 +71,53 @@ func clearChannelInfo(channel *model.Channel) {
 	}
 }
 
+func channelIDsFromChannels(channels []*model.Channel) []int {
+	ids := make([]int, 0, len(channels))
+	for _, channel := range channels {
+		if channel != nil && channel.Id > 0 {
+			ids = append(ids, channel.Id)
+		}
+	}
+	return ids
+}
+
+func closeActiveChannelWebSockets(channelIDs []int) {
+	service.CloseActiveWebSocketsForChannels(channelIDs, service.ChannelDisabledCloseReason)
+}
+
+func hasEnabledMultiKey(channel *model.Channel) bool {
+	if channel == nil || !channel.ChannelInfo.IsMultiKey {
+		return true
+	}
+	keys := channel.GetKeys()
+	if len(keys) == 0 {
+		return false
+	}
+	for i := range keys {
+		if channel.ChannelInfo.MultiKeyStatusList == nil {
+			return true
+		}
+		if status, ok := channel.ChannelInfo.MultiKeyStatusList[i]; !ok || status == common.ChannelStatusEnabled {
+			return true
+		}
+	}
+	return false
+}
+
+func disableMultiKeyChannelIfUnavailable(channel *model.Channel) bool {
+	if channel == nil || !channel.ChannelInfo.IsMultiKey || hasEnabledMultiKey(channel) {
+		return false
+	}
+	if channel.Status != common.ChannelStatusEnabled {
+		return true
+	}
+	if !model.UpdateChannelStatus(channel.Id, "", common.ChannelStatusManuallyDisabled, "All keys are disabled") {
+		return false
+	}
+	channel.Status = common.ChannelStatusManuallyDisabled
+	return true
+}
+
 func applyChannelStatusFilter(query *gorm.DB, statusFilter int) *gorm.DB {
 	if statusFilter == common.ChannelStatusEnabled {
 		return query.Where("status = ?", common.ChannelStatusEnabled)
@@ -740,6 +787,7 @@ func DeleteChannel(c *gin.Context) {
 		"id":   id,
 		"name": channelName,
 	})
+	closeActiveChannelWebSockets([]int{id})
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",
@@ -748,6 +796,13 @@ func DeleteChannel(c *gin.Context) {
 }
 
 func DeleteDisabledChannel(c *gin.Context) {
+	var ids []int
+	if err := model.DB.Model(&model.Channel{}).
+		Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).
+		Pluck("id", &ids).Error; err != nil {
+		common.ApiError(c, err)
+		return
+	}
 	rows, err := model.DeleteDisabledChannel()
 	if err != nil {
 		common.ApiError(c, err)
@@ -760,6 +815,7 @@ func DeleteDisabledChannel(c *gin.Context) {
 	recordManageAudit(c, "channel.delete_disabled", map[string]interface{}{
 		"count": rows,
 	})
+	closeActiveChannelWebSockets(ids)
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",
@@ -790,6 +846,12 @@ func DisableTagChannels(c *gin.Context) {
 		})
 		return
 	}
+	channels, err := model.GetChannelsByTag(channelTag.Tag, false, false)
+	if err != nil {
+		common.ApiError(c, err)
+		return
+	}
+	ids := channelIDsFromChannels(channels)
 	err = model.DisableChannelByTag(channelTag.Tag)
 	if err != nil {
 		common.ApiError(c, err)
@@ -799,6 +861,7 @@ func DisableTagChannels(c *gin.Context) {
 	recordManageAudit(c, "channel.tag_disable", map[string]interface{}{
 		"tag": channelTag.Tag,
 	})
+	closeActiveChannelWebSockets(ids)
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",
@@ -919,6 +982,7 @@ func DeleteChannelBatch(c *gin.Context) {
 	recordManageAudit(c, "channel.delete_batch", map[string]interface{}{
 		"count": deletedCount,
 	})
+	closeActiveChannelWebSockets(channelBatch.Ids)
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",
@@ -1138,6 +1202,9 @@ func UpdateChannelStatus(c *gin.Context) {
 	changed := model.UpdateChannelStatus(id, "", req.Status, "manual operation")
 	if changed {
 		model.InitChannelCache()
+		if req.Status != common.ChannelStatusEnabled {
+			closeActiveChannelWebSockets([]int{id})
+		}
 	}
 	recordManageAudit(c, "channel.status_update", map[string]interface{}{
 		"id":      id,
@@ -1165,6 +1232,9 @@ func BatchUpdateChannelStatus(c *gin.Context) {
 	}
 	if changedCount > 0 {
 		model.InitChannelCache()
+		if req.Status != common.ChannelStatusEnabled {
+			closeActiveChannelWebSockets(req.Ids)
+		}
 	}
 	recordManageAudit(c, "channel.status_update_batch", map[string]interface{}{
 		"count":  changedCount,
@@ -1686,7 +1756,11 @@ func ManageMultiKeys(c *gin.Context) {
 			return
 		}
 
+		shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
 		model.InitChannelCache()
+		if shouldCloseWebSocket {
+			closeActiveChannelWebSockets([]int{channel.Id})
+		}
 		c.JSON(http.StatusOK, gin.H{
 			"success": true,
 			"message": "密钥已禁用",
@@ -1799,7 +1873,11 @@ func ManageMultiKeys(c *gin.Context) {
 			return
 		}
 
+		shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
 		model.InitChannelCache()
+		if shouldCloseWebSocket {
+			closeActiveChannelWebSockets([]int{channel.Id})
+		}
 		c.JSON(http.StatusOK, gin.H{
 			"success": true,
 			"message": fmt.Sprintf("已禁用 %d 个密钥", disabledCount),
@@ -1879,7 +1957,11 @@ func ManageMultiKeys(c *gin.Context) {
 			return
 		}
 
+		shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
 		model.InitChannelCache()
+		if shouldCloseWebSocket {
+			closeActiveChannelWebSockets([]int{channel.Id})
+		}
 		c.JSON(http.StatusOK, gin.H{
 			"success": true,
 			"message": "密钥已删除",
@@ -1947,7 +2029,11 @@ func ManageMultiKeys(c *gin.Context) {
 			return
 		}
 
+		shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
 		model.InitChannelCache()
+		if shouldCloseWebSocket {
+			closeActiveChannelWebSockets([]int{channel.Id})
+		}
 		c.JSON(http.StatusOK, gin.H{
 			"success": true,
 			"message": fmt.Sprintf("已删除 %d 个自动禁用的密钥", deletedCount),