feat: sync deployed v1.0.0-rc.38-custom.42
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
cneicy committed
1 parent
aa6c3f50f4
commit
ceddfd3d6b
828 files changed
+122253
-13702
No files matched your search
@@ -0,0 +1,250 @@
|
||||
package wsmanager
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/go-redis/redis/v8"
|
||||
)
|
||||
|
||||
const (
|
||||
KindRealtime = "realtime"
|
||||
KindResponses = "responses"
|
||||
|
||||
defaultCloseReason = "channel disabled or deleted"
|
||||
redisChannel = "new-api:wsmanager:channel-close"
|
||||
)
|
||||
|
||||
type entry struct {
|
||||
id uint64
|
||||
channelID int
|
||||
kind string
|
||||
close func(reason string)
|
||||
}
|
||||
|
||||
type closeEvent struct {
|
||||
ChannelIDs []int `json:"channel_ids"`
|
||||
Reason string `json:"reason"`
|
||||
Origin string `json:"origin"`
|
||||
}
|
||||
|
||||
var (
|
||||
mu sync.Mutex
|
||||
nextID uint64
|
||||
registry = map[int]map[uint64]*entry{}
|
||||
|
||||
originOnce sync.Once
|
||||
originID string
|
||||
|
||||
subscriberOnce sync.Once
|
||||
)
|
||||
|
||||
func Register(channelID int, kind string, close func(reason string)) func() {
|
||||
if channelID <= 0 || close == nil {
|
||||
return func() {}
|
||||
}
|
||||
var closeOnce sync.Once
|
||||
safeClose := func(reason string) {
|
||||
closeOnce.Do(func() {
|
||||
close(reason)
|
||||
})
|
||||
}
|
||||
id := atomic.AddUint64(&nextID, 1)
|
||||
e := &entry{
|
||||
id: id,
|
||||
channelID: channelID,
|
||||
kind: kind,
|
||||
close: safeClose,
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
if registry[channelID] == nil {
|
||||
registry[channelID] = map[uint64]*entry{}
|
||||
}
|
||||
registry[channelID][id] = e
|
||||
mu.Unlock()
|
||||
|
||||
var once sync.Once
|
||||
return func() {
|
||||
once.Do(func() {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
entries := registry[channelID]
|
||||
if entries == nil {
|
||||
return
|
||||
}
|
||||
delete(entries, id)
|
||||
if len(entries) == 0 {
|
||||
delete(registry, channelID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func CloseChannel(channelID int, reason string) int {
|
||||
return CloseChannels([]int{channelID}, reason)
|
||||
}
|
||||
|
||||
func CloseChannels(channelIDs []int, reason string) int {
|
||||
reason = normalizeReason(reason)
|
||||
entries := takeEntries(channelIDs)
|
||||
for _, e := range entries {
|
||||
e.close(reason)
|
||||
}
|
||||
if len(entries) > 0 {
|
||||
common.SysLog(fmt.Sprintf("closed %d active websocket connection(s), channels=%v, kinds=%v, reason=%s", len(entries), entryChannelIDs(entries), entryKindCounts(entries), reason))
|
||||
}
|
||||
return len(entries)
|
||||
}
|
||||
|
||||
func CloseChannelsAndBroadcast(channelIDs []int, reason string) int {
|
||||
count := CloseChannels(channelIDs, reason)
|
||||
if err := PublishCloseChannels(context.Background(), channelIDs, reason); err != nil {
|
||||
common.SysLog(fmt.Sprintf("failed to publish websocket close event: %v", err))
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func StartSubscriber(ctx context.Context) {
|
||||
rdb := common.RDB
|
||||
if !common.RedisEnabled || rdb == nil {
|
||||
return
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
subscriberOnce.Do(func() {
|
||||
go func() {
|
||||
pubsub := rdb.Subscribe(ctx, channelCloseTopic(rdb.Options().DB))
|
||||
defer pubsub.Close()
|
||||
receiveChannelCloseEvents(ctx, pubsub.Channel(), getOriginID())
|
||||
}()
|
||||
})
|
||||
}
|
||||
|
||||
func PublishCloseChannels(ctx context.Context, channelIDs []int, reason string) error {
|
||||
rdb := common.RDB
|
||||
if !common.RedisEnabled || rdb == nil {
|
||||
return nil
|
||||
}
|
||||
ids := uniqueChannelIDs(channelIDs)
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
payload, err := common.Marshal(closeEvent{
|
||||
ChannelIDs: ids,
|
||||
Reason: normalizeReason(reason),
|
||||
Origin: getOriginID(),
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return rdb.Publish(ctx, channelCloseTopic(rdb.Options().DB), payload).Err()
|
||||
}
|
||||
|
||||
func channelCloseTopic(database int) string {
|
||||
// Redis Pub/Sub crosses logical databases, unlike the cache keys. Include
|
||||
// the selected database so independent deployments cannot close each other.
|
||||
return fmt.Sprintf("%s:db:%d", redisChannel, database)
|
||||
}
|
||||
|
||||
func receiveChannelCloseEvents(ctx context.Context, ch <-chan *redis.Message, origin string) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case msg, ok := <-ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var event closeEvent
|
||||
if err := common.Unmarshal([]byte(msg.Payload), &event); err != nil {
|
||||
common.SysLog(fmt.Sprintf("failed to unmarshal websocket close event: %v", err))
|
||||
continue
|
||||
}
|
||||
if event.Origin == origin {
|
||||
continue
|
||||
}
|
||||
CloseChannels(event.ChannelIDs, event.Reason)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func takeEntries(channelIDs []int) []*entry {
|
||||
ids := uniqueChannelIDs(channelIDs)
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
var entries []*entry
|
||||
for _, channelID := range ids {
|
||||
for _, e := range registry[channelID] {
|
||||
entries = append(entries, e)
|
||||
}
|
||||
delete(registry, channelID)
|
||||
}
|
||||
return entries
|
||||
}
|
||||
|
||||
func entryChannelIDs(entries []*entry) []int {
|
||||
ids := make([]int, 0, len(entries))
|
||||
for _, e := range entries {
|
||||
if e != nil {
|
||||
ids = append(ids, e.channelID)
|
||||
}
|
||||
}
|
||||
return uniqueChannelIDs(ids)
|
||||
}
|
||||
|
||||
func entryKindCounts(entries []*entry) map[string]int {
|
||||
counts := make(map[string]int)
|
||||
for _, e := range entries {
|
||||
if e == nil {
|
||||
continue
|
||||
}
|
||||
counts[e.kind]++
|
||||
}
|
||||
return counts
|
||||
}
|
||||
|
||||
func uniqueChannelIDs(channelIDs []int) []int {
|
||||
seen := make(map[int]struct{}, len(channelIDs))
|
||||
ids := make([]int, 0, len(channelIDs))
|
||||
for _, id := range channelIDs {
|
||||
if id <= 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func normalizeReason(reason string) string {
|
||||
if reason == "" {
|
||||
return defaultCloseReason
|
||||
}
|
||||
return reason
|
||||
}
|
||||
|
||||
func getOriginID() string {
|
||||
originOnce.Do(func() {
|
||||
name := common.NodeName
|
||||
if name == "" {
|
||||
name = "node"
|
||||
}
|
||||
originID = fmt.Sprintf("%s-%d-%d", name, os.Getpid(), time.Now().UnixNano())
|
||||
})
|
||||
return originID
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
package wsmanager
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func resetRegistryForTest() {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
registry = map[int]map[uint64]*entry{}
|
||||
nextID = 0
|
||||
}
|
||||
|
||||
func TestCloseChannelClosesRegisteredConnectionsOnce(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
|
||||
var mu sync.Mutex
|
||||
calls := 0
|
||||
Register(10, KindRealtime, func(reason string) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
require.Equal(t, "test reason", reason, "close callback should receive the provided reason")
|
||||
calls++
|
||||
})
|
||||
Register(10, KindResponses, func(reason string) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
calls++
|
||||
})
|
||||
|
||||
require.Equal(t, 2, CloseChannel(10, "test reason"), "CloseChannel should close every registered connection for the channel")
|
||||
require.Equal(t, 0, CloseChannel(10, "test reason"), "CloseChannel should not close already removed connections")
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
assert.Equal(t, 2, calls, "all registered close callbacks should be called once")
|
||||
}
|
||||
|
||||
func TestCloseChannelDoesNotCloseOtherChannels(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
|
||||
calls := map[int]int{}
|
||||
Register(10, KindRealtime, func(reason string) {
|
||||
calls[10]++
|
||||
})
|
||||
Register(20, KindRealtime, func(reason string) {
|
||||
calls[20]++
|
||||
})
|
||||
|
||||
require.Equal(t, 1, CloseChannel(10, "test"), "CloseChannel should only close the requested channel")
|
||||
assert.Equal(t, 1, calls[10], "requested channel callback should run")
|
||||
assert.Equal(t, 0, calls[20], "other channel callback should not run")
|
||||
}
|
||||
|
||||
func TestUnregisterPreventsClose(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
|
||||
calls := 0
|
||||
unregister := Register(10, KindRealtime, func(reason string) {
|
||||
calls++
|
||||
})
|
||||
unregister()
|
||||
|
||||
require.Equal(t, 0, CloseChannel(10, "test"), "unregistered connections should not be closed")
|
||||
assert.Equal(t, 0, calls, "unregistered close callback should not run")
|
||||
}
|
||||
|
||||
func TestRegisteredCloseIsIdempotent(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
|
||||
calls := 0
|
||||
Register(10, KindRealtime, func(reason string) {
|
||||
calls++
|
||||
})
|
||||
|
||||
mu.Lock()
|
||||
var registered *entry
|
||||
for _, e := range registry[10] {
|
||||
registered = e
|
||||
}
|
||||
mu.Unlock()
|
||||
require.NotNil(t, registered, "registered entry should exist")
|
||||
|
||||
registered.close("test")
|
||||
registered.close("test")
|
||||
assert.Equal(t, 1, calls, "registered close callback should be idempotent")
|
||||
}
|
||||
|
||||
func TestPublishCloseChannelsNoopsWhenRedisDisabled(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
|
||||
oldEnabled := common.RedisEnabled
|
||||
oldRDB := common.RDB
|
||||
common.RedisEnabled = false
|
||||
common.RDB = nil
|
||||
defer func() {
|
||||
common.RedisEnabled = oldEnabled
|
||||
common.RDB = oldRDB
|
||||
}()
|
||||
|
||||
require.NoError(t, PublishCloseChannels(context.Background(), []int{10}, "test"), "publishing should no-op without Redis")
|
||||
}
|
||||
|
||||
func TestRedisChannelCloseEventsStayWithinDatabase(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
previousEnabled, previousClient := common.RedisEnabled, common.RDB
|
||||
t.Cleanup(func() {
|
||||
common.RedisEnabled, common.RDB = previousEnabled, previousClient
|
||||
resetRegistryForTest()
|
||||
})
|
||||
address := os.Getenv("TEST_WS_MANAGER_REDIS_ADDR")
|
||||
if address == "" {
|
||||
address = miniredis.RunT(t).Addr()
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
t.Cleanup(cancel)
|
||||
localPublisher := redis.NewClient(&redis.Options{Addr: address, DB: 0})
|
||||
localSubscriber := redis.NewClient(&redis.Options{Addr: address, DB: 0})
|
||||
otherPublisher := redis.NewClient(&redis.Options{Addr: address, DB: 1})
|
||||
for _, client := range []*redis.Client{localPublisher, localSubscriber, otherPublisher} {
|
||||
require.NoError(t, client.Ping(ctx).Err())
|
||||
t.Cleanup(func() { require.NoError(t, client.Close()) })
|
||||
}
|
||||
if os.Getenv("TEST_WS_MANAGER_REDIS_ADDR") != "" {
|
||||
info, err := localPublisher.Info(ctx, "server").Result()
|
||||
require.NoError(t, err)
|
||||
for line := range strings.SplitSeq(info, "\r\n") {
|
||||
if strings.HasPrefix(line, "redis_version:") {
|
||||
t.Log(line)
|
||||
}
|
||||
}
|
||||
}
|
||||
pubsub := localSubscriber.Subscribe(ctx, channelCloseTopic(localSubscriber.Options().DB))
|
||||
t.Cleanup(func() { require.NoError(t, pubsub.Close()) })
|
||||
_, err := pubsub.Receive(ctx)
|
||||
require.NoError(t, err, "wait for subscription acknowledgement before publishing")
|
||||
otherPubsub := otherPublisher.Subscribe(ctx, channelCloseTopic(otherPublisher.Options().DB))
|
||||
t.Cleanup(func() { require.NoError(t, otherPubsub.Close()) })
|
||||
_, err = otherPubsub.Receive(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
closed := make(chan int, 2)
|
||||
Register(10, KindRealtime, func(string) { closed <- 10 })
|
||||
Register(20, KindResponses, func(string) { closed <- 20 })
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
receiveChannelCloseEvents(ctx, pubsub.Channel(), "receiving-node")
|
||||
}()
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
<-done
|
||||
})
|
||||
|
||||
common.RedisEnabled, common.RDB = true, otherPublisher
|
||||
require.NoError(t, PublishCloseChannels(ctx, []int{10}, "other database"))
|
||||
otherEvent, err := otherPubsub.ReceiveMessage(ctx)
|
||||
require.NoError(t, err)
|
||||
var event closeEvent
|
||||
require.NoError(t, common.Unmarshal([]byte(otherEvent.Payload), &event))
|
||||
assert.Equal(t, []int{10}, event.ChannelIDs)
|
||||
assert.Equal(t, "other database", event.Reason)
|
||||
// The old global topic must not affect a namespaced receiver either.
|
||||
require.NoError(t, localPublisher.Publish(ctx, redisChannel, otherEvent.Payload).Err())
|
||||
|
||||
common.RDB = localPublisher
|
||||
require.NoError(t, PublishCloseChannels(ctx, []int{20}, "same database"))
|
||||
select {
|
||||
case channelID := <-closed:
|
||||
assert.Equal(t, 20, channelID, "only the same-database broadcast may close a local connection")
|
||||
case <-ctx.Done():
|
||||
t.Fatal("same-database broadcast did not close its connection")
|
||||
}
|
||||
// Redis processes the preceding publications before this successful one,
|
||||
// so the retained registration proves isolation without sleeps or polling.
|
||||
assert.Equal(t, 1, CloseChannel(10, "test cleanup"))
|
||||
}
|
||||
Reference in new issue
Block a user