Add per-room bandwidth limits.

cmd/proxy/proxy_server.go

@@ -133,10 +133,11 @@ type ProxyServer struct {
 	stopped atomic.Bool
 	load    atomic.Uint64
 
-	maxIncoming     api.AtomicBandwidth
-	currentIncoming api.AtomicBandwidth
-	maxOutgoing     api.AtomicBandwidth
-	currentOutgoing api.AtomicBandwidth
+	maxIncoming       api.AtomicBandwidth
+	currentIncoming   api.AtomicBandwidth
+	maxOutgoing       api.AtomicBandwidth
+	currentOutgoing   api.AtomicBandwidth
+	currentBandwidths atomic.Pointer[map[string]*sfu.ClientBandwidthInfo]
 
 	shutdownChannel   chan struct{}
 	shutdownScheduled atomic.Bool
@@ -514,26 +515,32 @@ func (s *ProxyServer) newLoadEvent(load uint64, incoming api.Bandwidth, outgoing
 }
 
 func (s *ProxyServer) updateLoad() {
-	load, incoming, outgoing := s.GetClientsLoad()
+	load, incoming, outgoing, bandwidths := s.GetClientsLoad()
 	oldLoad := s.load.Swap(load)
 	oldIncoming := s.currentIncoming.Swap(incoming)
 	oldOutgoing := s.currentOutgoing.Swap(outgoing)
-	if oldLoad == load && oldIncoming == incoming && oldOutgoing == outgoing {
+	if len(bandwidths) == 0 {
+		s.currentBandwidths.Store(nil)
+	} else {
+		s.currentBandwidths.Store(&bandwidths)
+	}
+	if oldLoad == load && oldIncoming == incoming && oldOutgoing == outgoing && len(bandwidths) == 0 {
 		return
 	}
 
 	statsLoadCurrent.Set(float64(load))
-	s.sendLoadToAll(load, incoming, outgoing)
+	s.sendLoadToAll(load, incoming, outgoing, bandwidths)
 }
 
-func (s *ProxyServer) sendLoadToAll(load uint64, incoming api.Bandwidth, outgoing api.Bandwidth) {
+func (s *ProxyServer) sendLoadToAll(load uint64, incoming api.Bandwidth, outgoing api.Bandwidth, bandwidths map[string]*sfu.ClientBandwidthInfo) {
 	if s.shutdownScheduled.Load() {
 		// Server is scheduled to shutdown, no need to update clients with current load.
 		return
 	}
 
-	msg := s.newLoadEvent(load, incoming, outgoing)
+	loadEvent := s.newLoadEvent(load, incoming, outgoing)
 	s.IterateSessions(func(session *ProxySession) {
+		msg := session.updateLoadEvent(loadEvent, bandwidths)
 		session.sendMessage(msg)
 	})
 }
@@ -638,7 +645,11 @@ func (s *ProxyServer) loadConfig(config *goconf.ConfigFile, fromReload bool) err
 	oldOutgoing := s.maxOutgoing.Swap(maxOutgoing)
 	if fromReload && (oldIncoming != maxIncoming || oldOutgoing != maxOutgoing) {
 		// Notify sessions about updated load / bandwidth usage.
-		go s.sendLoadToAll(s.load.Load(), s.currentIncoming.Load(), s.currentOutgoing.Load())
+		var bandwidths map[string]*sfu.ClientBandwidthInfo
+		if bw := s.currentBandwidths.Load(); bw != nil {
+			bandwidths = *bw
+		}
+		go s.sendLoadToAll(s.load.Load(), s.currentIncoming.Load(), s.currentOutgoing.Load(), bandwidths)
 	}
 
 	return nil
@@ -734,7 +745,12 @@ func (s *ProxyServer) onMcuDisconnected() {
 }
 
 func (s *ProxyServer) sendCurrentLoad(session *ProxySession) {
+	var bandwidths map[string]*sfu.ClientBandwidthInfo
+	if bw := s.currentBandwidths.Load(); bw != nil {
+		bandwidths = *bw
+	}
 	msg := s.newLoadEvent(s.load.Load(), s.currentIncoming.Load(), s.currentOutgoing.Load())
+	msg = session.updateLoadEvent(msg, bandwidths)
 	session.sendMessage(msg)
 }
 
@@ -1200,6 +1216,37 @@ func (s *ProxyServer) processCommand(ctx context.Context, client *ProxyClient, s
 			client.Close(context.Background())
 		}()
 
+		response := &proxy.ServerMessage{
+			Id:   message.Id,
+			Type: "command",
+			Command: &proxy.CommandServerMessage{
+				Id: cmd.ClientId,
+			},
+		}
+		session.sendMessage(response)
+	case "update-bandwidth":
+		client := s.GetClient(cmd.ClientId)
+		if client == nil {
+			session.sendMessage(message.NewErrorServerMessage(UnknownClient))
+			return
+		}
+
+		clientWithBandwidth, ok := client.(sfu.ClientWithBandwidth)
+		if !ok {
+			session.sendMessage(message.NewErrorServerMessage(UnknownClient))
+			return
+		}
+
+		if cmd.Bandwidth != 0 {
+			ctx2, cancel := context.WithTimeout(ctx, s.mcuTimeout)
+			defer cancel()
+
+			if err := clientWithBandwidth.SetBandwidth(ctx2, cmd.Bandwidth); err != nil {
+				session.sendMessage(message.NewWrappedErrorServerMessage(err))
+				return
+			}
+		}
+
 		response := &proxy.ServerMessage{
 			Id:   message.Id,
 			Type: "command",
@@ -1582,14 +1629,18 @@ func (s *ProxyServer) HasClients() bool {
 	return len(s.clients) > 0
 }
 
-func (s *ProxyServer) GetClientsLoad() (load uint64, incoming api.Bandwidth, outgoing api.Bandwidth) {
+func (s *ProxyServer) GetClientsLoad() (load uint64, incoming api.Bandwidth, outgoing api.Bandwidth, bandwidths map[string]*sfu.ClientBandwidthInfo) {
 	s.clientsLock.RLock()
 	defer s.clientsLock.RUnlock()
 
-	for _, c := range s.clients {
+	for id, c := range s.clients {
 		// Use "current" bandwidth usage if supported.
 		if bw, ok := c.(sfu.ClientWithBandwidth); ok {
 			if bandwidth := bw.Bandwidth(); bandwidth != nil {
+				if bandwidths == nil {
+					bandwidths = make(map[string]*sfu.ClientBandwidthInfo)
+				}
+				bandwidths[id] = bandwidth
 				incoming += bandwidth.Received
 				outgoing += bandwidth.Sent
 				continue

cmd/proxy/proxy_server_test.go

@@ -459,13 +459,16 @@ func (p *TestMCUPublisher) UnpublishRemote(ctx context.Context, remoteId api.Pub
 
 type PublisherTestMCU struct {
 	TestMCU
+
+	publisher atomic.Pointer[TestPublisherWithBandwidth]
 }
 
 type TestPublisherWithBandwidth struct {
 	TestMCUPublisher
 
-	t         *testing.T
-	bandwidth *sfu.ClientBandwidthInfo
+	t            *testing.T
+	bandwidth    *sfu.ClientBandwidthInfo
+	bandwidthSet atomic.Bool
 }
 
 func (p *TestPublisherWithBandwidth) Bandwidth() *sfu.ClientBandwidthInfo {
@@ -491,6 +494,16 @@ func (p *TestPublisherWithBandwidth) SendMessage(ctx context.Context, message *a
 	}
 }
 
+func (p *TestPublisherWithBandwidth) SetBandwidth(ctx context.Context, bandwidth api.Bandwidth) error {
+	assert.EqualValues(p.t, 20000, bandwidth)
+	p.bandwidthSet.Store(true)
+	return nil
+}
+
+func (m *PublisherTestMCU) GetPublisher() *TestPublisherWithBandwidth {
+	return m.publisher.Load()
+}
+
 func (m *PublisherTestMCU) NewPublisher(ctx context.Context, listener sfu.Listener, id api.PublicSessionId, sid string, streamType sfu.StreamType, settings sfu.NewPublisherSettings, initiator sfu.Initiator) (sfu.Publisher, error) {
 	publisher := &TestPublisherWithBandwidth{
 		TestMCUPublisher: TestMCUPublisher{
@@ -505,6 +518,10 @@ func (m *PublisherTestMCU) NewPublisher(ctx context.Context, listener sfu.Listen
 			Received: api.BandwidthFromBytes(2000),
 		},
 	}
+	if !m.publisher.CompareAndSwap(nil, publisher) {
+		return nil, errors.New("only one publisher supported")
+	}
+
 	return publisher, nil
 }
 
@@ -571,6 +588,8 @@ func TestProxyPublisherBandwidth(t *testing.T) {
 		assert.Equal(clientId, proxyServer.GetClientId(publisher))
 	}
 
+	publisher := mcu.GetPublisher()
+	require.NotNil(publisher)
 	proxyServer.updateLoad()
 
 	if message, err := client.RunUntilMessage(ctx); assert.NoError(err) {
@@ -584,9 +603,34 @@ func TestProxyPublisherBandwidth(t *testing.T) {
 					assert.InEpsilon(10, *bw.Outgoing, 0.0001)
 				}
 			}
+			if assert.Len(message.Event.ClientBandwidths, 1) {
+				if bw := message.Event.ClientBandwidths; assert.NotNil(bw[clientId], "expected %s, got %+v", bw) {
+					assert.EqualValues(8000, bw[clientId].Sent)
+					assert.EqualValues(16000, bw[clientId].Received)
+				}
+			}
 		}
 	}
 
+	require.NoError(client.WriteJSON(&proxy.ClientMessage{
+		Id:   "3456",
+		Type: "command",
+		Command: &proxy.CommandClientMessage{
+			Type:      "update-bandwidth",
+			ClientId:  clientId,
+			Bandwidth: api.BandwidthFromBits(20000),
+		},
+	}))
+
+	if message, err := client.RunUntilMessage(ctx); assert.NoError(err) {
+		assert.Equal("3456", message.Id)
+		if err := checkMessageType(message, "command"); assert.NoError(err) {
+			assert.Equal(clientId, message.Command.Id)
+		}
+	}
+
+	assert.True(publisher.bandwidthSet.Load(), "should have set bandwidth")
+
 	require.NoError(client.WriteJSON(&proxy.ClientMessage{
 		Id:   "3456",
 		Type: "payload",

cmd/proxy/proxy_session.go

@@ -484,3 +484,77 @@ func (s *ProxySession) OnRemotePublisherDeleted(publisherId api.PublicSessionId)
 		}
 	}
 }
+
+func cloneLoadMessageWithoutClientBandwidths(msg *proxy.ServerMessage) *proxy.ServerMessage {
+	return &proxy.ServerMessage{
+		Id:   msg.Id,
+		Type: msg.Type,
+		Event: &proxy.EventServerMessage{
+			Type:      msg.Event.Type,
+			ClientId:  msg.Event.ClientId,
+			Load:      msg.Event.Load,
+			Sid:       msg.Event.Sid,
+			Bandwidth: msg.Event.Bandwidth,
+		},
+	}
+}
+
+func (s *ProxySession) updateLoadEventPublishers(msg *proxy.ServerMessage, bandwidths map[string]*sfu.ClientBandwidthInfo, needClone bool) (result *proxy.ServerMessage, cloned bool) {
+	s.publishersLock.Lock()
+	defer s.publishersLock.Unlock()
+
+	result = msg
+	for id := range s.publishers {
+		if bw, found := bandwidths[id]; found {
+			if needClone {
+				result = cloneLoadMessageWithoutClientBandwidths(msg)
+				needClone = false
+				cloned = true
+			}
+
+			if result.Event.ClientBandwidths == nil {
+				result.Event.ClientBandwidths = make(map[string]proxy.EventServerBandwidth)
+			}
+			result.Event.ClientBandwidths[id] = proxy.EventServerBandwidth{
+				Received: bw.Received,
+				Sent:     bw.Sent,
+			}
+		}
+	}
+	return
+}
+
+func (s *ProxySession) updateLoadEventSubscribers(msg *proxy.ServerMessage, bandwidths map[string]*sfu.ClientBandwidthInfo, needClone bool) (result *proxy.ServerMessage, cloned bool) {
+	s.subscribersLock.Lock()
+	defer s.subscribersLock.Unlock()
+
+	result = msg
+	for id := range s.subscribers {
+		if bw, found := bandwidths[id]; found {
+			if needClone {
+				result = cloneLoadMessageWithoutClientBandwidths(msg)
+				needClone = false
+				cloned = true
+			}
+
+			if result.Event.ClientBandwidths == nil {
+				result.Event.ClientBandwidths = make(map[string]proxy.EventServerBandwidth)
+			}
+			result.Event.ClientBandwidths[id] = proxy.EventServerBandwidth{
+				Received: bw.Received,
+				Sent:     bw.Sent,
+			}
+		}
+	}
+	return
+}
+
+func (s *ProxySession) updateLoadEvent(msg *proxy.ServerMessage, bandwidths map[string]*sfu.ClientBandwidthInfo) *proxy.ServerMessage {
+	if len(bandwidths) == 0 {
+		return msg
+	}
+
+	msg, cloned := s.updateLoadEventPublishers(msg, bandwidths, true)
+	msg, _ = s.updateLoadEventSubscribers(msg, bandwidths, !cloned)
+	return msg
+}

etcd/api.go

@@ -69,6 +69,10 @@ type BackendInformationEtcd struct {
 	MaxStreamBitrate api.Bandwidth `json:"maxstreambitrate,omitempty"`
 	MaxScreenBitrate api.Bandwidth `json:"maxscreenbitrate,omitempty"`
 
+	BandwidthPerRoom      api.Bandwidth `json:"bitrateperroom,omitempty"`
+	MinPublisherBandwidth api.Bandwidth `json:"minpublisherbitrate,omitempty"`
+	MaxPublisherBandwidth api.Bandwidth `json:"maxpublisherbitrate,omitempty"`
+
 	SessionLimit uint64 `json:"sessionlimit,omitempty"`
 }