Merge commit '6bd7fac875d9e9009915053d8a590abb372c5679' into feature/livekit-upgrade-1.13.1

This commit is contained in:
2026-06-25 23:54:33 +09:00
291 changed files with 37631 additions and 15381 deletions
+74 -42
View File
@@ -23,6 +23,7 @@ import (
"go.uber.org/zap/zapcore"
"google.golang.org/protobuf/proto"
"github.com/livekit/livekit-server/pkg/config"
"github.com/livekit/livekit-server/pkg/utils"
"github.com/livekit/protocol/auth"
"github.com/livekit/protocol/livekit"
@@ -94,7 +95,7 @@ type NullMessageSource struct {
func NewNullMessageSource(connID livekit.ConnectionID) *NullMessageSource {
return &NullMessageSource{
connID: connID,
msgChan: make(chan proto.Message, 0),
msgChan: make(chan proto.Message),
}
}
@@ -152,8 +153,13 @@ type StartParticipantSignalResults struct {
type MessageRouter interface {
// CreateRoom starts an rtc room
CreateRoom(ctx context.Context, req *livekit.CreateRoomRequest) (res *livekit.Room, err error)
// StartParticipantSignal participant signal connection is ready to start
StartParticipantSignal(ctx context.Context, roomName livekit.RoomName, pi ParticipantInit) (res StartParticipantSignalResults, err error)
StartParticipantSignal(
ctx context.Context,
roomName livekit.RoomName,
pi ParticipantInit,
) (res StartParticipantSignalResults, err error)
}
func CreateRouter(
@@ -162,8 +168,9 @@ func CreateRouter(
signalClient SignalClient,
roomManagerClient RoomManagerClient,
kps rpc.KeepalivePubSub,
nodeStatsConfig config.NodeStatsConfig,
) Router {
lr := NewLocalRouter(node, signalClient, roomManagerClient)
lr := NewLocalRouter(node, signalClient, roomManagerClient, nodeStatsConfig)
if rc != nil {
return NewRedisRouter(lr, rc, kps)
@@ -177,19 +184,24 @@ func CreateRouter(
// ------------------------------------------------
type ParticipantInit struct {
Identity livekit.ParticipantIdentity
Name livekit.ParticipantName
Reconnect bool
ReconnectReason livekit.ReconnectReason
AutoSubscribe bool
Client *livekit.ClientInfo
Grants *auth.ClaimGrants
Region string
AdaptiveStream bool
ID livekit.ParticipantID
SubscriberAllowPause *bool
DisableICELite bool
CreateRoom *livekit.CreateRoomRequest
Identity livekit.ParticipantIdentity
Name livekit.ParticipantName
Reconnect bool
ReconnectReason livekit.ReconnectReason
AutoSubscribe bool
AutoSubscribeDataTrack *bool
Client *livekit.ClientInfo
Grants *auth.ClaimGrants
Region string
AdaptiveStream bool
ID livekit.ParticipantID
SubscriberAllowPause *bool
DisableICELite bool
CreateRoom *livekit.CreateRoomRequest
AddTrackRequests []*livekit.AddTrackRequest
PublisherOffer *livekit.SessionDescription
SyncState *livekit.SyncState
UseSinglePeerConnection bool
}
func (pi *ParticipantInit) MarshalLogObject(e zapcore.ObjectEncoder) error {
@@ -209,6 +221,7 @@ func (pi *ParticipantInit) MarshalLogObject(e zapcore.ObjectEncoder) error {
logBoolPtr("Reconnect", &pi.Reconnect)
e.AddString("ReconnectReason", pi.ReconnectReason.String())
logBoolPtr("AutoSubscribe", &pi.AutoSubscribe)
logBoolPtr("AutoSubscribeDataTrack", pi.AutoSubscribeDataTrack)
e.AddObject("Client", logger.Proto(utils.ClientInfoWithoutAddress(pi.Client)))
e.AddObject("Grants", pi.Grants)
e.AddString("Region", pi.Region)
@@ -217,6 +230,10 @@ func (pi *ParticipantInit) MarshalLogObject(e zapcore.ObjectEncoder) error {
logBoolPtr("SubscriberAllowPause", pi.SubscriberAllowPause)
logBoolPtr("DisableICELite", &pi.DisableICELite)
e.AddObject("CreateRoom", logger.Proto(pi.CreateRoom))
e.AddArray("AddTrackRequests", logger.ProtoSlice(pi.AddTrackRequests))
e.AddObject("PublisherOffer", logger.Proto(pi.PublisherOffer))
e.AddObject("SyncState", logger.Proto(pi.SyncState))
logBoolPtr("UseSinglePeerConnection", &pi.UseSinglePeerConnection)
return nil
}
@@ -227,20 +244,27 @@ func (pi *ParticipantInit) ToStartSession(roomName livekit.RoomName, connectionI
}
ss := &livekit.StartSession{
RoomName: string(roomName),
Identity: string(pi.Identity),
Name: string(pi.Name),
// connection id is to allow the RTC node to identify where to route the message back to
ConnectionId: string(connectionID),
Reconnect: pi.Reconnect,
ReconnectReason: pi.ReconnectReason,
AutoSubscribe: pi.AutoSubscribe,
Client: pi.Client,
GrantsJson: string(claims),
AdaptiveStream: pi.AdaptiveStream,
ParticipantId: string(pi.ID),
DisableIceLite: pi.DisableICELite,
CreateRoom: pi.CreateRoom,
RoomName: string(roomName),
Identity: string(pi.Identity),
Name: string(pi.Name),
ConnectionId: string(connectionID),
Reconnect: pi.Reconnect,
ReconnectReason: pi.ReconnectReason,
AutoSubscribe: pi.AutoSubscribe,
Client: pi.Client,
GrantsJson: string(claims),
AdaptiveStream: pi.AdaptiveStream,
ParticipantId: string(pi.ID),
DisableIceLite: pi.DisableICELite,
CreateRoom: pi.CreateRoom,
AddTrackRequests: pi.AddTrackRequests,
PublisherOffer: pi.PublisherOffer,
SyncState: pi.SyncState,
UseSinglePeerConnection: pi.UseSinglePeerConnection,
}
if pi.AutoSubscribeDataTrack != nil {
autoSubscribeDataTrack := *pi.AutoSubscribeDataTrack
ss.AutoSubscribeDataTrack = &autoSubscribeDataTrack
}
if pi.SubscriberAllowPause != nil {
subscriberAllowPause := *pi.SubscriberAllowPause
@@ -257,18 +281,26 @@ func ParticipantInitFromStartSession(ss *livekit.StartSession, region string) (*
}
pi := &ParticipantInit{
Identity: livekit.ParticipantIdentity(ss.Identity),
Name: livekit.ParticipantName(ss.Name),
Reconnect: ss.Reconnect,
ReconnectReason: ss.ReconnectReason,
Client: ss.Client,
AutoSubscribe: ss.AutoSubscribe,
Grants: claims,
Region: region,
AdaptiveStream: ss.AdaptiveStream,
ID: livekit.ParticipantID(ss.ParticipantId),
DisableICELite: ss.DisableIceLite,
CreateRoom: ss.CreateRoom,
Identity: livekit.ParticipantIdentity(ss.Identity),
Name: livekit.ParticipantName(ss.Name),
Reconnect: ss.Reconnect,
ReconnectReason: ss.ReconnectReason,
Client: ss.Client,
AutoSubscribe: ss.AutoSubscribe,
Grants: claims,
Region: region,
AdaptiveStream: ss.AdaptiveStream,
ID: livekit.ParticipantID(ss.ParticipantId),
DisableICELite: ss.DisableIceLite,
CreateRoom: ss.CreateRoom,
AddTrackRequests: ss.AddTrackRequests,
PublisherOffer: ss.PublisherOffer,
SyncState: ss.SyncState,
UseSinglePeerConnection: ss.UseSinglePeerConnection,
}
if ss.AutoSubscribeDataTrack != nil {
autoSubscribeDataTrack := *ss.AutoSubscribeDataTrack
pi.AutoSubscribeDataTrack = &autoSubscribeDataTrack
}
if ss.SubscriberAllowPause != nil {
subscriberAllowPause := *ss.SubscriberAllowPause
+7 -5
View File
@@ -16,11 +16,11 @@ package routing
import (
"context"
"sync"
"time"
"go.uber.org/atomic"
"github.com/livekit/livekit-server/pkg/config"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
)
@@ -32,8 +32,8 @@ type LocalRouter struct {
currentNode LocalNode
signalClient SignalClient
roomManagerClient RoomManagerClient
nodeStatsConfig config.NodeStatsConfig
lock sync.RWMutex
// channels for each participant
requestChannels map[string]*MessageChannel
responseChannels map[string]*MessageChannel
@@ -44,11 +44,13 @@ func NewLocalRouter(
currentNode LocalNode,
signalClient SignalClient,
roomManagerClient RoomManagerClient,
nodeStatsConfig config.NodeStatsConfig,
) *LocalRouter {
return &LocalRouter{
currentNode: currentNode,
signalClient: signalClient,
roomManagerClient: roomManagerClient,
nodeStatsConfig: nodeStatsConfig,
requestChannels: make(map[string]*MessageChannel),
responseChannels: make(map[string]*MessageChannel),
}
@@ -106,7 +108,8 @@ func (r *LocalRouter) StartParticipantSignal(ctx context.Context, roomName livek
func (r *LocalRouter) StartParticipantSignalWithNodeID(ctx context.Context, roomName livekit.RoomName, pi ParticipantInit, nodeID livekit.NodeID) (res StartParticipantSignalResults, err error) {
connectionID, reqSink, resSource, err := r.signalClient.StartParticipantSignal(ctx, roomName, pi, nodeID)
if err != nil {
logger.Errorw("could not handle new participant", err,
logger.Errorw(
"could not handle new participant", err,
"room", roomName,
"participant", pi.Identity,
"connID", connectionID,
@@ -146,8 +149,7 @@ func (r *LocalRouter) statsWorker() {
if !r.isStarted.Load() {
return
}
// update every 10 seconds
<-time.After(statsUpdateInterval)
<-time.After(r.nodeStatsConfig.StatsUpdateInterval)
r.currentNode.UpdateNodeStats()
}
}
@@ -38,7 +38,7 @@ func TestMessageChannel_WriteMessageClosed(t *testing.T) {
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 100; i++ {
for range 100 {
_ = m.WriteMessage(&livekit.SignalRequest{})
}
}()
+14 -17
View File
@@ -20,12 +20,10 @@ import (
"time"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
"github.com/livekit/protocol/utils"
"github.com/livekit/protocol/utils/guid"
"github.com/livekit/livekit-server/pkg/config"
"github.com/livekit/livekit-server/pkg/telemetry/prometheus"
)
type LocalNode interface {
@@ -45,30 +43,35 @@ type LocalNodeImpl struct {
lock sync.RWMutex
node *livekit.Node
// previous stats for computing averages
prevStats *livekit.NodeStats
nodeStats *NodeStats
}
func NewLocalNode(conf *config.Config) (*LocalNodeImpl, error) {
nodeID := guid.New(utils.NodePrefix)
if conf != nil && conf.RTC.NodeIP == "" {
if conf != nil && conf.RTC.NodeIP.IsEmpty() {
return nil, ErrIPNotSet
}
nowUnix := time.Now().Unix()
l := &LocalNodeImpl{
node: &livekit.Node{
Id: nodeID,
NumCpus: uint32(runtime.NumCPU()),
State: livekit.NodeState_SERVING,
Stats: &livekit.NodeStats{
StartedAt: time.Now().Unix(),
UpdatedAt: time.Now().Unix(),
StartedAt: nowUnix,
UpdatedAt: nowUnix,
},
},
}
var nsc *config.NodeStatsConfig
if conf != nil {
l.node.Ip = conf.RTC.NodeIP
l.node.Ip = conf.RTC.NodeIP.PrimaryIP()
l.node.Region = conf.Region
nsc = &conf.NodeStats
}
l.nodeStats = NewNodeStats(nsc, nowUnix)
return l, nil
}
@@ -138,18 +141,12 @@ func (l *LocalNodeImpl) UpdateNodeStats() bool {
l.lock.Lock()
defer l.lock.Unlock()
if l.prevStats == nil {
l.prevStats = l.node.Stats
}
updated, computedAvg, err := prometheus.GetUpdatedNodeStats(l.node.Stats, l.prevStats)
stats, err := l.nodeStats.UpdateAndGetNodeStats()
if err != nil {
logger.Errorw("could not update node stats", err)
return false
}
l.node.Stats = updated
if computedAvg {
l.prevStats = updated
}
l.node.Stats = stats
return true
}
+82
View File
@@ -0,0 +1,82 @@
// Copyright 2023 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package routing
import (
"sync"
"time"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
"github.com/livekit/livekit-server/pkg/config"
"github.com/livekit/livekit-server/pkg/telemetry/prometheus"
)
type NodeStats struct {
config config.NodeStatsConfig
startedAt int64
lock sync.Mutex
statsHistory []*livekit.NodeStats
statsHistoryWritePtr int
}
func NewNodeStats(conf *config.NodeStatsConfig, startedAt int64) *NodeStats {
n := &NodeStats{
startedAt: startedAt,
}
n.UpdateConfig(conf)
return n
}
func (n *NodeStats) UpdateConfig(conf *config.NodeStatsConfig) {
n.lock.Lock()
defer n.lock.Unlock()
if conf == nil {
conf = &config.DefaultNodeStatsConfig
}
n.config = *conf
// set up stats history to be able to measure different rate windows
var maxInterval time.Duration
for _, rateInterval := range conf.StatsRateMeasurementIntervals {
if rateInterval > maxInterval {
maxInterval = rateInterval
}
}
n.statsHistory = make([]*livekit.NodeStats, (maxInterval+conf.StatsUpdateInterval-1)/conf.StatsUpdateInterval)
n.statsHistoryWritePtr = 0
}
func (n *NodeStats) UpdateAndGetNodeStats() (*livekit.NodeStats, error) {
n.lock.Lock()
defer n.lock.Unlock()
stats, err := prometheus.GetNodeStats(
n.startedAt,
append(n.statsHistory[n.statsHistoryWritePtr:], n.statsHistory[0:n.statsHistoryWritePtr]...),
n.config.StatsRateMeasurementIntervals,
)
if err != nil {
logger.Errorw("could not update node stats", err)
return nil, err
}
n.statsHistory[n.statsHistoryWritePtr] = stats
n.statsHistoryWritePtr = (n.statsHistoryWritePtr + 1) % len(n.statsHistory)
return stats, nil
}
+3 -8
View File
@@ -33,11 +33,6 @@ import (
)
const (
// expire participant mappings after a day
participantMappingTTL = 24 * time.Hour
statsUpdateInterval = 2 * time.Second
statsMaxDelaySeconds = float64(30)
// hash of node_id => Node proto
NodesKey = "nodes"
@@ -209,11 +204,11 @@ func (r *RedisRouter) statsWorker() {
for r.ctx.Err() == nil {
// update periodically
select {
case <-time.After(statsUpdateInterval):
case <-time.After(r.nodeStatsConfig.StatsUpdateInterval):
r.kps.PublishPing(r.ctx, r.currentNode.NodeID(), &rpc.KeepalivePing{Timestamp: time.Now().Unix()})
delaySeconds := r.currentNode.SecondsSinceNodeStatsUpdate()
if delaySeconds > statsMaxDelaySeconds {
if delaySeconds > r.nodeStatsConfig.StatsMaxDelay.Seconds() {
if !goroutineDumped {
goroutineDumped = true
buf := bytes.NewBuffer(nil)
@@ -240,7 +235,7 @@ func (r *RedisRouter) keepaliveWorker(startedChan chan error) {
close(startedChan)
for ping := range pings.Channel() {
if time.Since(time.Unix(ping.Timestamp, 0)) > statsUpdateInterval {
if time.Since(time.Unix(ping.Timestamp, 0)) > r.nodeStatsConfig.StatsUpdateInterval {
logger.Infow("keep alive too old, skipping", "timestamp", ping.Timestamp)
continue
}
@@ -243,14 +243,6 @@ func (fake *FakeMessageSink) WriteMessageReturnsOnCall(i int, result1 error) {
func (fake *FakeMessageSink) Invocations() map[string][][]interface{} {
fake.invocationsMutex.RLock()
defer fake.invocationsMutex.RUnlock()
fake.closeMutex.RLock()
defer fake.closeMutex.RUnlock()
fake.connectionIDMutex.RLock()
defer fake.connectionIDMutex.RUnlock()
fake.isClosedMutex.RLock()
defer fake.isClosedMutex.RUnlock()
fake.writeMessageMutex.RLock()
defer fake.writeMessageMutex.RUnlock()
copiedInvocations := map[string][][]interface{}{}
for key, value := range fake.invocations {
copiedInvocations[key] = value
@@ -234,14 +234,6 @@ func (fake *FakeMessageSource) ReadChanReturnsOnCall(i int, result1 <-chan proto
func (fake *FakeMessageSource) Invocations() map[string][][]interface{} {
fake.invocationsMutex.RLock()
defer fake.invocationsMutex.RUnlock()
fake.closeMutex.RLock()
defer fake.closeMutex.RUnlock()
fake.connectionIDMutex.RLock()
defer fake.connectionIDMutex.RUnlock()
fake.isClosedMutex.RLock()
defer fake.isClosedMutex.RUnlock()
fake.readChanMutex.RLock()
defer fake.readChanMutex.RUnlock()
copiedInvocations := map[string][][]interface{}{}
for key, value := range fake.invocations {
copiedInvocations[key] = value
@@ -129,10 +129,6 @@ func (fake *FakeRoomManagerClient) CreateRoomReturnsOnCall(i int, result1 *livek
func (fake *FakeRoomManagerClient) Invocations() map[string][][]interface{} {
fake.invocationsMutex.RLock()
defer fake.invocationsMutex.RUnlock()
fake.closeMutex.RLock()
defer fake.closeMutex.RUnlock()
fake.createRoomMutex.RLock()
defer fake.createRoomMutex.RUnlock()
copiedInvocations := map[string][][]interface{}{}
for key, value := range fake.invocations {
copiedInvocations[key] = value
@@ -845,32 +845,6 @@ func (fake *FakeRouter) UnregisterNodeReturnsOnCall(i int, result1 error) {
func (fake *FakeRouter) Invocations() map[string][][]interface{} {
fake.invocationsMutex.RLock()
defer fake.invocationsMutex.RUnlock()
fake.clearRoomStateMutex.RLock()
defer fake.clearRoomStateMutex.RUnlock()
fake.createRoomMutex.RLock()
defer fake.createRoomMutex.RUnlock()
fake.drainMutex.RLock()
defer fake.drainMutex.RUnlock()
fake.getNodeForRoomMutex.RLock()
defer fake.getNodeForRoomMutex.RUnlock()
fake.getRegionMutex.RLock()
defer fake.getRegionMutex.RUnlock()
fake.listNodesMutex.RLock()
defer fake.listNodesMutex.RUnlock()
fake.registerNodeMutex.RLock()
defer fake.registerNodeMutex.RUnlock()
fake.removeDeadNodesMutex.RLock()
defer fake.removeDeadNodesMutex.RUnlock()
fake.setNodeForRoomMutex.RLock()
defer fake.setNodeForRoomMutex.RUnlock()
fake.startMutex.RLock()
defer fake.startMutex.RUnlock()
fake.startParticipantSignalMutex.RLock()
defer fake.startParticipantSignalMutex.RUnlock()
fake.stopMutex.RLock()
defer fake.stopMutex.RUnlock()
fake.unregisterNodeMutex.RLock()
defer fake.unregisterNodeMutex.RUnlock()
copiedInvocations := map[string][][]interface{}{}
for key, value := range fake.invocations {
copiedInvocations[key] = value
@@ -173,10 +173,6 @@ func (fake *FakeSignalClient) StartParticipantSignalReturnsOnCall(i int, result1
func (fake *FakeSignalClient) Invocations() map[string][][]interface{} {
fake.invocationsMutex.RLock()
defer fake.invocationsMutex.RUnlock()
fake.activeCountMutex.RLock()
defer fake.activeCountMutex.RUnlock()
fake.startParticipantSignalMutex.RLock()
defer fake.startParticipantSignalMutex.RUnlock()
copiedInvocations := map[string][][]interface{}{}
for key, value := range fake.invocations {
copiedInvocations[key] = value
+3 -2
View File
@@ -20,7 +20,8 @@ import (
// AnySelector selects any available node with no limitations
type AnySelector struct {
SortBy string
SortBy string
Algorithm string
}
func (s *AnySelector) SelectNode(nodes []*livekit.Node) (*livekit.Node, error) {
@@ -29,5 +30,5 @@ func (s *AnySelector) SelectNode(nodes []*livekit.Node) (*livekit.Node, error) {
return nil, ErrNoAvailableNodes
}
return SelectSortedNode(nodes, s.SortBy)
return SelectSortedNode(nodes, s.SortBy, s.Algorithm)
}
@@ -0,0 +1,287 @@
// Copyright 2023 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package selector
import (
"testing"
"time"
"github.com/stretchr/testify/require"
"github.com/livekit/protocol/livekit"
)
func createTestNode(id string, cpuLoad float32, numRooms int32, numClients int32, state livekit.NodeState) *livekit.Node {
return &livekit.Node{
Id: id,
State: state,
Stats: &livekit.NodeStats{
UpdatedAt: time.Now().Unix() - 1, // Recent update to be considered available
CpuLoad: cpuLoad,
NumRooms: numRooms,
NumClients: numClients,
NumCpus: 4,
LoadAvgLast1Min: cpuLoad * 4, // Simulate system load
},
}
}
func TestAnySelector_SelectNode_TwoChoice(t *testing.T) {
tests := []struct {
name string
sortBy string
algorithm string
nodes []*livekit.Node
wantErr string
expected string
notExpected string
}{
{
name: "successful selection with cpuload sorting",
sortBy: "cpuload",
algorithm: "twochoice",
nodes: []*livekit.Node{
createTestNode("node1", 0.8, 5, 10, livekit.NodeState_SERVING),
createTestNode("node2", 0.3, 2, 5, livekit.NodeState_SERVING),
createTestNode("node3", 0.6, 3, 8, livekit.NodeState_SERVING),
createTestNode("node4", 0.9, 6, 12, livekit.NodeState_SERVING),
},
wantErr: "",
expected: "", // Not determinstic selection, so no specific expected node
notExpected: "node4", // Node with highest load should not be selected
},
{
name: "successful selection with rooms sorting",
sortBy: "rooms",
algorithm: "twochoice",
nodes: []*livekit.Node{
createTestNode("node1", 0.5, 8, 15, livekit.NodeState_SERVING),
createTestNode("node2", 0.4, 2, 5, livekit.NodeState_SERVING),
createTestNode("node3", 0.6, 12, 20, livekit.NodeState_SERVING),
},
wantErr: "",
expected: "", // Not determinstic selection, so no specific expected node
notExpected: "node3", // Node with highest room count should not be selected
},
{
name: "successful selection with clients sorting",
sortBy: "clients",
algorithm: "twochoice",
nodes: []*livekit.Node{
createTestNode("node1", 0.5, 3, 25, livekit.NodeState_SERVING),
createTestNode("node2", 0.4, 2, 5, livekit.NodeState_SERVING),
createTestNode("node3", 0.6, 4, 30, livekit.NodeState_SERVING),
},
wantErr: "",
expected: "", // Not determinstic selection, so no specific expected node
notExpected: "node3", // Node with highest clients should not be selected
},
{
name: "empty nodes list",
sortBy: "cpuload",
algorithm: "twochoice",
nodes: []*livekit.Node{},
wantErr: "could not find any available nodes",
},
{
name: "no available nodes - all unavailable",
sortBy: "cpuload",
algorithm: "twochoice",
nodes: []*livekit.Node{
{
Id: "node1",
State: livekit.NodeState_SERVING,
Stats: &livekit.NodeStats{
UpdatedAt: time.Now().Unix() - 10, // Too old
CpuLoad: 0.3,
},
},
},
wantErr: "could not find any available nodes",
},
{
name: "no available nodes - not serving",
sortBy: "cpuload",
algorithm: "twochoice",
nodes: []*livekit.Node{
{
Id: "node1",
State: livekit.NodeState_SHUTTING_DOWN,
Stats: &livekit.NodeStats{
UpdatedAt: time.Now().Unix() - 1,
CpuLoad: 0.3,
},
},
},
wantErr: "could not find any available nodes",
},
{
name: "single available node",
sortBy: "cpuload",
algorithm: "twochoice",
nodes: []*livekit.Node{
createTestNode("node1", 0.5, 3, 10, livekit.NodeState_SERVING),
},
wantErr: "",
expected: "node1", // Should select the only available node
notExpected: "", // No other nodes to compare against
},
{
name: "two available nodes",
sortBy: "cpuload",
algorithm: "twochoice",
nodes: []*livekit.Node{
createTestNode("node1", 0.8, 5, 15, livekit.NodeState_SERVING),
createTestNode("node2", 0.3, 2, 5, livekit.NodeState_SERVING),
},
wantErr: "",
expected: "node2", // Should select the node with lower load
notExpected: "node1", // Should not select the node with higher load
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
selector := &AnySelector{
SortBy: tt.sortBy,
Algorithm: tt.algorithm,
}
node, err := selector.SelectNode(tt.nodes)
if tt.wantErr != "" {
require.Error(t, err)
require.Contains(t, err.Error(), tt.wantErr)
require.Nil(t, node)
} else {
require.NoError(t, err)
require.NotNil(t, node)
require.NotEmpty(t, node.Id)
// Verify the selected node is one of the available nodes
found := false
availableNodes := GetAvailableNodes(tt.nodes)
for _, availableNode := range availableNodes {
if availableNode.Id == node.Id {
found = true
break
}
}
require.True(t, found, "Selected node should be one of the available nodes")
if tt.expected != "" {
require.Equal(t, tt.expected, node.Id, "Selected node should match expected")
}
if tt.notExpected != "" {
require.NotEqual(t, tt.notExpected, node.Id, "Selected node should not match not expected")
}
}
})
}
}
func TestAnySelector_SelectNode_TwoChoice_Probabilistic_Behavior(t *testing.T) {
// Test that two-choice algorithm favors nodes with lower metrics
// This test runs multiple iterations to increase confidence in the probabilistic behavior
selector := &AnySelector{
SortBy: "cpuload",
Algorithm: "twochoice",
}
// Create nodes where node2 has significantly lower CPU load
nodes := []*livekit.Node{
createTestNode("node1", 0.95, 10, 20, livekit.NodeState_SERVING), // Very high load
createTestNode("node2", 0.1, 1, 2, livekit.NodeState_SERVING), // Low load
createTestNode("node3", 0.5, 9, 18, livekit.NodeState_SERVING), // Medium load
createTestNode("node4", 0.85, 8, 16, livekit.NodeState_SERVING), // High load
}
// Run multiple selections and count how often the low-load node is selected
iterations := 1000
lowLoadSelections := 0
higestLoadSelections := 0
for range iterations {
node, err := selector.SelectNode(nodes)
require.NoError(t, err)
require.NotNil(t, node)
if node.Id == "node2" {
lowLoadSelections++
}
if node.Id == "node1" {
higestLoadSelections++
}
}
// The low-load node should be selected more often than pure random (25%)
// Due to the two-choice algorithm favoring the better node
selectionRate := float64(lowLoadSelections) / float64(iterations)
require.Greater(t, selectionRate, 0.4, "Two-choice algorithm should favor the low-load node more than random selection")
require.Equal(t, higestLoadSelections, 0, "Two-choice algorithm should never favor the highest load node")
}
func TestAnySelector_SelectNode_InvalidParameters(t *testing.T) {
nodes := []*livekit.Node{
createTestNode("node1", 0.5, 3, 10, livekit.NodeState_SERVING),
}
tests := []struct {
name string
sortBy string
algorithm string
wantErr string
}{
{
name: "empty sortBy",
sortBy: "",
algorithm: "twochoice",
wantErr: "sort by option cannot be blank",
},
{
name: "empty algorithm",
sortBy: "cpuload",
algorithm: "",
wantErr: "node selector algorithm option cannot be blank",
},
{
name: "unknown sortBy",
sortBy: "invalid",
algorithm: "twochoice",
wantErr: "unknown sort by option",
},
{
name: "unknown algorithm",
sortBy: "cpuload",
algorithm: "invalid",
wantErr: "unknown node selector algorithm option",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
selector := &AnySelector{
SortBy: tt.sortBy,
Algorithm: tt.algorithm,
}
node, err := selector.SelectNode(nodes)
require.Error(t, err)
require.Contains(t, err.Error(), tt.wantErr)
require.Nil(t, node)
})
}
}
@@ -23,6 +23,7 @@ import (
type CPULoadSelector struct {
CPULoadLimit float32
SortBy string
Algorithm string
}
func (s *CPULoadSelector) filterNodes(nodes []*livekit.Node) ([]*livekit.Node, error) {
@@ -50,5 +51,5 @@ func (s *CPULoadSelector) SelectNode(nodes []*livekit.Node) (*livekit.Node, erro
return nil, err
}
return SelectSortedNode(nodes, s.SortBy)
return SelectSortedNode(nodes, s.SortBy, s.Algorithm)
}
@@ -25,7 +25,7 @@ import (
)
func TestCPULoadSelector_SelectNode(t *testing.T) {
sel := selector.CPULoadSelector{CPULoadLimit: 0.8, SortBy: "random"}
sel := selector.CPULoadSelector{CPULoadLimit: 0.8, SortBy: "random", Algorithm: "lowest"}
var nodes []*livekit.Node
_, err := sel.SelectNode(nodes)
@@ -39,7 +39,7 @@ func TestCPULoadSelector_SelectNode(t *testing.T) {
// Select a node with low load when available
nodes = []*livekit.Node{nodeLoadLow, nodeLoadHigh}
for i := 0; i < 5; i++ {
for range 5 {
node, err := sel.SelectNode(nodes)
if err != nil {
t.Error(err)
@@ -21,5 +21,7 @@ var (
ErrCurrentRegionNotSet = errors.New("current region cannot be blank")
ErrCurrentRegionUnknownLatLon = errors.New("unknown lat and lon for the current region")
ErrSortByNotSet = errors.New("sort by option cannot be blank")
ErrAlgorithmNotSet = errors.New("node selector algorithm option cannot be blank")
ErrSortByUnknown = errors.New("unknown sort by option")
ErrAlgorithmUnknown = errors.New("unknown node selector algorithm option")
)
@@ -37,19 +37,21 @@ func CreateNodeSelector(conf *config.Config) (NodeSelector, error) {
}
switch kind {
case "any":
return &AnySelector{conf.NodeSelector.SortBy}, nil
return &AnySelector{conf.NodeSelector.SortBy, conf.NodeSelector.Algorithm}, nil
case "cpuload":
return &CPULoadSelector{
CPULoadLimit: conf.NodeSelector.CPULoadLimit,
SortBy: conf.NodeSelector.SortBy,
SortBy: conf.NodeSelector.SortBy,
Algorithm: conf.NodeSelector.Algorithm,
}, nil
case "sysload":
return &SystemLoadSelector{
SysloadLimit: conf.NodeSelector.SysloadLimit,
SortBy: conf.NodeSelector.SortBy,
SortBy: conf.NodeSelector.SortBy,
Algorithm: conf.NodeSelector.Algorithm,
}, nil
case "regionaware":
s, err := NewRegionAwareSelector(conf.Region, conf.NodeSelector.Regions, conf.NodeSelector.SortBy)
s, err := NewRegionAwareSelector(conf.Region, conf.NodeSelector.Regions, conf.NodeSelector.SortBy, conf.NodeSelector.Algorithm)
if err != nil {
return nil, err
}
@@ -57,7 +59,7 @@ func CreateNodeSelector(conf *config.Config) (NodeSelector, error) {
return s, nil
case "random":
logger.Warnw("random node selector is deprecated, please switch to \"any\" or another selector", nil)
return &AnySelector{conf.NodeSelector.SortBy}, nil
return &AnySelector{conf.NodeSelector.SortBy, conf.NodeSelector.Algorithm}, nil
default:
return nil, ErrUnsupportedSelector
}
@@ -29,9 +29,10 @@ type RegionAwareSelector struct {
regionDistances map[string]float64
regions []config.RegionConfig
SortBy string
Algorithm string
}
func NewRegionAwareSelector(currentRegion string, regions []config.RegionConfig, sortBy string) (*RegionAwareSelector, error) {
func NewRegionAwareSelector(currentRegion string, regions []config.RegionConfig, sortBy string, algorithm string) (*RegionAwareSelector, error) {
if currentRegion == "" {
return nil, ErrCurrentRegionNotSet
}
@@ -41,6 +42,7 @@ func NewRegionAwareSelector(currentRegion string, regions []config.RegionConfig,
regionDistances: make(map[string]float64),
regions: regions,
SortBy: sortBy,
Algorithm: algorithm,
}
var currentRC *config.RegionConfig
@@ -94,7 +96,7 @@ func (s *RegionAwareSelector) SelectNode(nodes []*livekit.Node) (*livekit.Node,
nodes = nearestNodes
}
return SelectSortedNode(nodes, s.SortBy)
return SelectSortedNode(nodes, s.SortBy, s.Algorithm)
}
// haversine(θ) function
@@ -34,6 +34,7 @@ const (
regionEast = "us-east"
regionSeattle = "seattle"
sortBy = "random"
algorithm = "lowest"
)
func TestRegionAwareRouting(t *testing.T) {
@@ -58,7 +59,7 @@ func TestRegionAwareRouting(t *testing.T) {
nodes := []*livekit.Node{
newTestNodeInRegion("", false),
}
s, err := selector.NewRegionAwareSelector(regionEast, nil, sortBy)
s, err := selector.NewRegionAwareSelector(regionEast, nil, sortBy, algorithm)
require.NoError(t, err)
node, err := s.SelectNode(nodes)
@@ -74,7 +75,7 @@ func TestRegionAwareRouting(t *testing.T) {
expectedNode,
newTestNodeInRegion(regionEast, false),
}
s, err := selector.NewRegionAwareSelector(regionEast, rc, sortBy)
s, err := selector.NewRegionAwareSelector(regionEast, rc, sortBy, algorithm)
require.NoError(t, err)
s.SysloadLimit = loadLimit
@@ -91,7 +92,7 @@ func TestRegionAwareRouting(t *testing.T) {
newTestNodeInRegion(regionWest, true),
newTestNodeInRegion(regionEast, false),
}
s, err := selector.NewRegionAwareSelector(regionEast, rc, sortBy)
s, err := selector.NewRegionAwareSelector(regionEast, rc, sortBy, algorithm)
require.NoError(t, err)
s.SysloadLimit = loadLimit
@@ -107,7 +108,7 @@ func TestRegionAwareRouting(t *testing.T) {
expectedNode,
newTestNodeInRegion(regionEast, true),
}
s, err := selector.NewRegionAwareSelector(regionSeattle, rc, sortBy)
s, err := selector.NewRegionAwareSelector(regionSeattle, rc, sortBy, algorithm)
require.NoError(t, err)
s.SysloadLimit = loadLimit
@@ -125,7 +126,7 @@ func TestRegionAwareRouting(t *testing.T) {
expectedNode,
expectedNode,
}
s, err := selector.NewRegionAwareSelector(regionSeattle, rc, sortBy)
s, err := selector.NewRegionAwareSelector(regionSeattle, rc, sortBy, algorithm)
require.NoError(t, err)
s.SysloadLimit = loadLimit
@@ -138,7 +139,7 @@ func TestRegionAwareRouting(t *testing.T) {
nodes := []*livekit.Node{
newTestNodeInRegion(regionWest, true),
}
s, err := selector.NewRegionAwareSelector(regionEast, rc, sortBy)
s, err := selector.NewRegionAwareSelector(regionEast, rc, sortBy, algorithm)
require.NoError(t, err)
node, err := s.SelectNode(nodes)
@@ -23,10 +23,10 @@ import (
)
func SortByTest(t *testing.T, sortBy string) {
sel := selector.SystemLoadSelector{SortBy: sortBy}
sel := selector.SystemLoadSelector{SortBy: sortBy, Algorithm: "lowest"}
nodes := []*livekit.Node{nodeLoadLow, nodeLoadMedium, nodeLoadHigh}
for i := 0; i < 5; i++ {
for range 5 {
node, err := sel.SelectNode(nodes)
if err != nil {
t.Error(err)
@@ -38,7 +38,7 @@ func SortByTest(t *testing.T, sortBy string) {
}
func TestSortByErrors(t *testing.T) {
sel := selector.SystemLoadSelector{}
sel := selector.SystemLoadSelector{Algorithm: "lowest"}
nodes := []*livekit.Node{nodeLoadLow, nodeLoadMedium, nodeLoadHigh}
// Test unset sort by option error
@@ -23,6 +23,7 @@ import (
type SystemLoadSelector struct {
SysloadLimit float32
SortBy string
Algorithm string
}
func (s *SystemLoadSelector) filterNodes(nodes []*livekit.Node) ([]*livekit.Node, error) {
@@ -49,5 +50,5 @@ func (s *SystemLoadSelector) SelectNode(nodes []*livekit.Node) (*livekit.Node, e
return nil, err
}
return SelectSortedNode(nodes, s.SortBy)
return SelectSortedNode(nodes, s.SortBy, s.Algorithm)
}
@@ -33,12 +33,16 @@ var (
NumCpus: 1,
CpuLoad: 0.1,
LoadAvgLast1Min: 0.0,
NumRooms: 1,
NumClients: 2,
NumTracksIn: 4,
NumTracksOut: 8,
BytesInPerSec: 1000,
BytesOutPerSec: 2000,
NumRooms: 1,
NumClients: 2,
NumTracksIn: 4,
NumTracksOut: 8,
Rates: []*livekit.NodeStatsRate{
{
BytesIn: 1000,
BytesOut: 2000,
},
},
},
}
@@ -49,12 +53,16 @@ var (
NumCpus: 1,
CpuLoad: 0.5,
LoadAvgLast1Min: 0.5,
NumRooms: 5,
NumClients: 10,
NumTracksIn: 20,
NumTracksOut: 200,
BytesInPerSec: 5000,
BytesOutPerSec: 10000,
NumRooms: 5,
NumClients: 10,
NumTracksIn: 20,
NumTracksOut: 200,
Rates: []*livekit.NodeStatsRate{
{
BytesIn: 5000,
BytesOut: 10000,
},
},
},
}
@@ -65,18 +73,22 @@ var (
NumCpus: 1,
CpuLoad: 0.99,
LoadAvgLast1Min: 2.0,
NumRooms: 10,
NumClients: 20,
NumTracksIn: 40,
NumTracksOut: 800,
BytesInPerSec: 10000,
BytesOutPerSec: 40000,
NumRooms: 10,
NumClients: 20,
NumTracksIn: 40,
NumTracksOut: 800,
Rates: []*livekit.NodeStatsRate{
{
BytesIn: 10000,
BytesOut: 40000,
},
},
},
}
)
func TestSystemLoadSelector_SelectNode(t *testing.T) {
sel := selector.SystemLoadSelector{SysloadLimit: 1.0, SortBy: "random"}
sel := selector.SystemLoadSelector{SysloadLimit: 1.0, SortBy: "random", Algorithm: "lowest"}
var nodes []*livekit.Node
_, err := sel.SelectNode(nodes)
@@ -90,7 +102,7 @@ func TestSystemLoadSelector_SelectNode(t *testing.T) {
// Select a node with low load when available
nodes = []*livekit.Node{nodeLoadLow, nodeLoadHigh}
for i := 0; i < 5; i++ {
for range 5 {
node, err := sel.SelectNode(nodes)
if err != nil {
t.Error(err)
+79 -15
View File
@@ -15,7 +15,8 @@
package selector
import (
"sort"
"math/rand/v2"
"slices"
"time"
"github.com/thoas/go-funk"
@@ -23,6 +24,7 @@ import (
"github.com/livekit/protocol/livekit"
"github.com/livekit/livekit-server/pkg/config"
"github.com/livekit/livekit-server/pkg/utils"
)
const AvailableSeconds = 5
@@ -62,54 +64,116 @@ func LimitsReached(limitConfig config.LimitConfig, nodeStats *livekit.NodeStats)
if limitConfig.NumTracks > 0 && limitConfig.NumTracks <= nodeStats.NumTracksIn+nodeStats.NumTracksOut {
return true
}
if limitConfig.BytesPerSec > 0 && limitConfig.BytesPerSec <= nodeStats.BytesInPerSec+nodeStats.BytesOutPerSec {
rate := &livekit.NodeStatsRate{}
if len(nodeStats.Rates) > 0 {
rate = nodeStats.Rates[0]
}
if limitConfig.BytesPerSec > 0 && limitConfig.BytesPerSec <= rate.BytesIn+rate.BytesOut {
return true
}
return false
}
func SelectSortedNode(nodes []*livekit.Node, sortBy string) (*livekit.Node, error) {
func SelectSortedNode(nodes []*livekit.Node, sortBy string, algorithm string) (*livekit.Node, error) {
if sortBy == "" {
return nil, ErrSortByNotSet
}
if algorithm == "" {
return nil, ErrAlgorithmNotSet
}
switch algorithm {
case "lowest": // examine all nodes and select the lowest based on sort criteria
return selectLowestSortedNode(nodes, sortBy)
case "twochoice": // randomly select two nodes and return the lowest based on sort criteria "Power of Two Random Choices"
return selectTwoChoiceSortedNode(nodes, sortBy)
default:
return nil, ErrAlgorithmUnknown
}
}
func selectTwoChoiceSortedNode(nodes []*livekit.Node, sortBy string) (*livekit.Node, error) {
if len(nodes) <= 2 {
return selectLowestSortedNode(nodes, sortBy)
}
// randomly select two nodes
node1, node2, err := selectTwoRandomNodes(nodes)
if err != nil {
return nil, err
}
// compare the two nodes based on the sort criteria
if node1 == nil || node2 == nil {
return nil, ErrNoAvailableNodes
}
selectedNode, err := selectLowestSortedNode([]*livekit.Node{node1, node2}, sortBy)
if err != nil {
return nil, err
}
return selectedNode, nil
}
func selectLowestSortedNode(nodes []*livekit.Node, sortBy string) (*livekit.Node, error) {
// Return a node based on what it should be sorted by for priority
switch sortBy {
case "random":
idx := funk.RandomInt(0, len(nodes))
return nodes[idx], nil
case "sysload":
sort.Slice(nodes, func(i, j int) bool {
return GetNodeSysload(nodes[i]) < GetNodeSysload(nodes[j])
slices.SortFunc(nodes, func(a, b *livekit.Node) int {
return utils.Signum(GetNodeSysload(a) - GetNodeSysload(b))
})
return nodes[0], nil
case "cpuload":
sort.Slice(nodes, func(i, j int) bool {
return nodes[i].Stats.CpuLoad < nodes[j].Stats.CpuLoad
slices.SortFunc(nodes, func(a, b *livekit.Node) int {
return utils.Signum(a.Stats.CpuLoad - b.Stats.CpuLoad)
})
return nodes[0], nil
case "rooms":
sort.Slice(nodes, func(i, j int) bool {
return nodes[i].Stats.NumRooms < nodes[j].Stats.NumRooms
slices.SortFunc(nodes, func(a, b *livekit.Node) int {
return utils.Signum(a.Stats.NumRooms - b.Stats.NumRooms)
})
return nodes[0], nil
case "clients":
sort.Slice(nodes, func(i, j int) bool {
return nodes[i].Stats.NumClients < nodes[j].Stats.NumClients
slices.SortFunc(nodes, func(a, b *livekit.Node) int {
return utils.Signum(a.Stats.NumClients - b.Stats.NumClients)
})
return nodes[0], nil
case "tracks":
sort.Slice(nodes, func(i, j int) bool {
return nodes[i].Stats.NumTracksIn+nodes[i].Stats.NumTracksOut < nodes[j].Stats.NumTracksIn+nodes[j].Stats.NumTracksOut
slices.SortFunc(nodes, func(a, b *livekit.Node) int {
return utils.Signum((a.Stats.NumTracksIn + a.Stats.NumTracksOut) - (b.Stats.NumTracksIn + b.Stats.NumTracksOut))
})
return nodes[0], nil
case "bytespersec":
sort.Slice(nodes, func(i, j int) bool {
return nodes[i].Stats.BytesInPerSec+nodes[i].Stats.BytesOutPerSec < nodes[j].Stats.BytesInPerSec+nodes[j].Stats.BytesOutPerSec
slices.SortFunc(nodes, func(a, b *livekit.Node) int {
ratea := &livekit.NodeStatsRate{}
if len(a.Stats.Rates) > 0 {
ratea = a.Stats.Rates[0]
}
rateb := &livekit.NodeStatsRate{}
if len(b.Stats.Rates) > 0 {
rateb = b.Stats.Rates[0]
}
return utils.Signum((ratea.BytesIn + ratea.BytesOut) - (rateb.BytesIn + rateb.BytesOut))
})
return nodes[0], nil
default:
return nil, ErrSortByUnknown
}
}
func selectTwoRandomNodes(nodes []*livekit.Node) (*livekit.Node, *livekit.Node, error) {
if len(nodes) < 2 {
return nil, nil, ErrNoAvailableNodes
}
shuffledIndices := rand.Perm(len(nodes))
return nodes[shuffledIndices[0]], nodes[shuffledIndices[1]], nil
}
+25 -8
View File
@@ -25,6 +25,7 @@ import (
"github.com/livekit/livekit-server/pkg/config"
"github.com/livekit/livekit-server/pkg/telemetry/prometheus"
"github.com/livekit/livekit-server/pkg/utils"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
"github.com/livekit/protocol/rpc"
@@ -50,7 +51,7 @@ type signalClient struct {
}
func NewSignalClient(nodeID livekit.NodeID, bus psrpc.MessageBus, config config.SignalRelayConfig) (SignalClient, error) {
c, err := rpc.NewTypedSignalClient(
client, err := rpc.NewTypedSignalClient(
nodeID,
bus,
middleware.WithClientMetrics(rpc.PSRPCMetricsObserver{}),
@@ -63,7 +64,7 @@ func NewSignalClient(nodeID livekit.NodeID, bus psrpc.MessageBus, config config.
return &signalClient{
nodeID: nodeID,
config: config,
client: c,
client: client,
}, nil
}
@@ -88,25 +89,27 @@ func (r *signalClient) StartParticipantSignal(
return
}
l := logger.GetLogger().WithValues(
l := utils.GetLogger(ctx).WithValues(
"room", roomName,
"reqNodeID", nodeID,
"participant", pi.Identity,
"connID", connectionID,
"participantInit", &pi,
"startSession", logger.Proto(ss),
)
l.Debugw("starting signal connection")
stream, err := r.client.RelaySignal(ctx, nodeID)
if err != nil {
prometheus.MessageCounter.WithLabelValues("signal", "failure").Add(1)
prometheus.RecordSignalRequestFailure()
return
}
err = stream.Send(&rpc.RelaySignalRequest{StartSession: ss})
if err != nil {
stream.Close(err)
prometheus.MessageCounter.WithLabelValues("signal", "failure").Add(1)
prometheus.RecordSignalRequestFailure()
return
}
@@ -130,6 +133,8 @@ func (r *signalClient) StartParticipantSignal(
resChan,
signalResponseMessageReader{},
r.config,
prometheus.RecordSignalResponseSuccess,
prometheus.RecordSignalResponseFailure,
)
l.Debugw("signal stream closed", "error", err)
@@ -139,6 +144,8 @@ func (r *signalClient) StartParticipantSignal(
return connectionID, sink, resChan, nil
}
// ------------------------------
type signalRequestMessageWriter struct{}
func (e signalRequestMessageWriter) Write(seq uint64, close bool, msgs []proto.Message) *rpc.RelaySignalRequest {
@@ -153,6 +160,8 @@ func (e signalRequestMessageWriter) Write(seq uint64, close bool, msgs []proto.M
return r
}
// -------------------------------
type signalResponseMessageReader struct{}
func (e signalResponseMessageReader) Read(rm *rpc.RelaySignalResponse) ([]proto.Message, error) {
@@ -163,6 +172,8 @@ func (e signalResponseMessageReader) Read(rm *rpc.RelaySignalResponse) ([]proto.
return msgs, nil
}
// -----------------------------------------
type RelaySignalMessage interface {
proto.Message
GetSeq() uint64
@@ -182,6 +193,8 @@ func CopySignalStreamToMessageChannel[SendType, RecvType RelaySignalMessage](
ch *MessageChannel,
reader SignalMessageReader[RecvType],
config config.SignalRelayConfig,
promSignalSuccess func(),
promSignalFailure func(),
) error {
r := &signalMessageReader[SendType, RecvType]{
reader: reader,
@@ -190,16 +203,16 @@ func CopySignalStreamToMessageChannel[SendType, RecvType RelaySignalMessage](
for msg := range stream.Channel() {
res, err := r.Read(msg)
if err != nil {
prometheus.MessageCounter.WithLabelValues("signal", "failure").Add(1)
promSignalFailure()
return err
}
for _, r := range res {
if err = ch.WriteMessage(r); err != nil {
prometheus.MessageCounter.WithLabelValues("signal", "failure").Add(1)
promSignalFailure()
return err
}
prometheus.MessageCounter.WithLabelValues("signal", "success").Add(1)
promSignalSuccess()
}
if msg.GetClose() {
@@ -209,6 +222,8 @@ func CopySignalStreamToMessageChannel[SendType, RecvType RelaySignalMessage](
return stream.Err()
}
// ----------------------------------------
type signalMessageReader[SendType, RecvType RelaySignalMessage] struct {
seq uint64
reader SignalMessageReader[RecvType]
@@ -236,6 +251,8 @@ func (r *signalMessageReader[SendType, RecvType]) Read(msg RecvType) ([]proto.Me
return res, nil
}
// ----------------------------------------
type SignalSinkParams[SendType, RecvType RelaySignalMessage] struct {
Stream psrpc.Stream[SendType, RecvType]
Logger logger.Logger