R3944Realms/srs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Codex: Remove global proxy load balancer.

Browse Source 
Create the origin load balancer in proxy bootstrap and pass it explicitly into protocol servers and system API handlers. Also keep load balancer implementations package-private and rename the default debugging origin helper.

Co-authored-by: chatgpt-codex-connector[bot] <199175422+chatgpt-codex-connector[bot]@users.noreply.github.com>
This commit is contained in:
winlin 2026-05-10 21:00:17 -04:00
parent 3b93ddfddf
commit f45bf30b46
11 changed files with 111 additions and 73 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
.gitignore vendored
View File

@ -47,4 +47,9 @@ cmake-build-debug
# For AI # For AI
/*personal* /*personal*
/support*
/*srs-consults*
/*workspace* /*workspace*
/skills/llm-switcher
/skills/*workspace*
/memory/202*.md

32
internal/bootstrap/proxy.go
View File

@ -65,8 +65,9 @@ func (b *proxyBootstrap) run(ctx context.Context) error {
// Start the Go pprof if enabled. // Start the Go pprof if enabled.
debug.HandleGoPprof(ctx, environment) debug.HandleGoPprof(ctx, environment)
// Initialize the load balancer. // Create and initialize the load balancer.
if err := b.initializeLoadBalancer(ctx, environment); err != nil { loadBalancer, err := b.initializeLoadBalancer(ctx, environment)
if err != nil {
return err return err
} }
@ -77,36 +78,37 @@ func (b *proxyBootstrap) run(ctx context.Context) error {
} }
// Start all servers and block until context is cancelled. // Start all servers and block until context is cancelled.
return b.startServers(ctx, environment, gracefulQuitTimeout) return b.startServers(ctx, environment, loadBalancer, gracefulQuitTimeout)
} }
// initializeLoadBalancer sets up the load balancer based on configuration. // initializeLoadBalancer sets up the load balancer based on configuration.
func (b *proxyBootstrap) initializeLoadBalancer(ctx context.Context, environment env.ProxyEnvironment) error { func (b *proxyBootstrap) initializeLoadBalancer(ctx context.Context, environment env.ProxyEnvironment) (lb.OriginLoadBalancer, error) {
var loadBalancer lb.OriginLoadBalancer
switch environment.LoadBalancerType() { switch environment.LoadBalancerType() {
case "redis": case "redis":
lb.SrsLoadBalancer = lb.NewRedisLoadBalancer(environment) loadBalancer = lb.NewRedisLoadBalancer(environment)
default: default:
lb.SrsLoadBalancer = lb.NewMemoryLoadBalancer(environment) loadBalancer = lb.NewMemoryLoadBalancer(environment)
} }
if err := lb.SrsLoadBalancer.Initialize(ctx); err != nil { if err := loadBalancer.Initialize(ctx); err != nil {
return errors.Wrapf(err, "initialize srs load balancer") return nil, errors.Wrapf(err, "initialize srs load balancer")
} }
return nil return loadBalancer, nil
} }
// startServers initializes and starts all protocol servers. // startServers initializes and starts all protocol servers.
func (b *proxyBootstrap) startServers(ctx context.Context, environment env.ProxyEnvironment, gracefulQuitTimeout time.Duration) error { func (b *proxyBootstrap) startServers(ctx context.Context, environment env.ProxyEnvironment, loadBalancer lb.OriginLoadBalancer, gracefulQuitTimeout time.Duration) error {
// Start the RTMP server. // Start the RTMP server.
rtmpProxyServer := proxy.NewRTMPProxyServer(environment) rtmpProxyServer := proxy.NewRTMPProxyServer(environment, loadBalancer)
if err := rtmpProxyServer.Run(ctx); err != nil { if err := rtmpProxyServer.Run(ctx); err != nil {
return errors.Wrapf(err, "rtmp server") return errors.Wrapf(err, "rtmp server")
} }
defer rtmpProxyServer.Close() defer rtmpProxyServer.Close()
// Start the WebRTC server. // Start the WebRTC server.
webRTCProxyServer := proxy.NewWebRTCProxyServer(environment) webRTCProxyServer := proxy.NewWebRTCProxyServer(environment, loadBalancer)
if err := webRTCProxyServer.Run(ctx); err != nil { if err := webRTCProxyServer.Run(ctx); err != nil {
return errors.Wrapf(err, "rtc server") return errors.Wrapf(err, "rtc server")
} }
@ -120,21 +122,21 @@ func (b *proxyBootstrap) startServers(ctx context.Context, environment env.Proxy
defer httpAPIProxyServer.Close() defer httpAPIProxyServer.Close()
// Start the SRT server. // Start the SRT server.
srsSRTProxyServer := proxy.NewSRSSRTProxyServer(environment) srsSRTProxyServer := proxy.NewSRSSRTProxyServer(environment, loadBalancer)
if err := srsSRTProxyServer.Run(ctx); err != nil { if err := srsSRTProxyServer.Run(ctx); err != nil {
return errors.Wrapf(err, "srt server") return errors.Wrapf(err, "srt server")
} }
defer srsSRTProxyServer.Close() defer srsSRTProxyServer.Close()
// Start the System API server. // Start the System API server.
systemAPI := proxy.NewSystemAPI(environment, gracefulQuitTimeout) systemAPI := proxy.NewSystemAPI(environment, loadBalancer, gracefulQuitTimeout)
if err := systemAPI.Run(ctx); err != nil { if err := systemAPI.Run(ctx); err != nil {
return errors.Wrapf(err, "system api server") return errors.Wrapf(err, "system api server")
} }
defer systemAPI.Close() defer systemAPI.Close()
// Start the HTTP web server. // Start the HTTP web server.
httpStreamProxyServer := proxy.NewHTTPStreamProxyServer(environment, gracefulQuitTimeout) httpStreamProxyServer := proxy.NewHTTPStreamProxyServer(environment, loadBalancer, gracefulQuitTimeout)
if err := httpStreamProxyServer.Run(ctx); err != nil { if err := httpStreamProxyServer.Run(ctx); err != nil {
return errors.Wrapf(err, "http server") return errors.Wrapf(err, "http server")
} }

4
internal/lb/debug.go
View File

@ -12,8 +12,8 @@ import (
"srsx/internal/logger" "srsx/internal/logger"
) )
// NewDefaultSRSForDebugging initialize the default SRS media server, for debugging only. // NewDefaultOriginServerForDebugging initializes the default origin server, for debugging only.
func NewDefaultSRSForDebugging(environment env.ProxyEnvironment) (*OriginServer, error) { func NewDefaultOriginServerForDebugging(environment env.ProxyEnvironment) (*OriginServer, error) {
if environment.DefaultBackendEnabled() != "on" { if environment.DefaultBackendEnabled() != "on" {
return nil, nil return nil, nil
} }

3
internal/lb/lb.go
View File

@ -126,6 +126,3 @@ type OriginLoadBalancer interface {
// Load the WebRTC streaming by ufrag, the ICE username. // Load the WebRTC streaming by ufrag, the ICE username.
LoadWebRTCByUfrag(ctx context.Context, ufrag string) (RTCConnection, error) LoadWebRTCByUfrag(ctx context.Context, ufrag string) (RTCConnection, error)
} }
// SrsLoadBalancer is the global SRS load balancer instance.
var SrsLoadBalancer OriginLoadBalancer

22
internal/lb/mem.go
View File

@ -15,8 +15,8 @@ import (
"srsx/internal/sync" "srsx/internal/sync"
) )
// MemoryLoadBalancer stores state in memory. // memoryLoadBalancer stores state in memory.
type MemoryLoadBalancer struct { type memoryLoadBalancer struct {
// The environment interface. // The environment interface.
environment env.ProxyEnvironment environment env.ProxyEnvironment
// All available SRS servers, key is server ID. // All available SRS servers, key is server ID.
@ -35,7 +35,7 @@ type MemoryLoadBalancer struct {
// NewMemoryLoadBalancer creates a new memory-based load balancer. // NewMemoryLoadBalancer creates a new memory-based load balancer.
func NewMemoryLoadBalancer(environment env.ProxyEnvironment) OriginLoadBalancer { func NewMemoryLoadBalancer(environment env.ProxyEnvironment) OriginLoadBalancer {
return &MemoryLoadBalancer{ return &memoryLoadBalancer{
environment: environment, environment: environment,
servers: sync.NewMap[string, *OriginServer](), servers: sync.NewMap[string, *OriginServer](),
picked: sync.NewMap[string, *OriginServer](), picked: sync.NewMap[string, *OriginServer](),
@ -46,8 +46,8 @@ func NewMemoryLoadBalancer(environment env.ProxyEnvironment) OriginLoadBalancer
} }
} }
func (v *MemoryLoadBalancer) Initialize(ctx context.Context) error { func (v *memoryLoadBalancer) Initialize(ctx context.Context) error {
server, err := NewDefaultSRSForDebugging(v.environment) server, err := NewDefaultOriginServerForDebugging(v.environment)
if err != nil { if err != nil {
return errors.Wrapf(err, "initialize default SRS") return errors.Wrapf(err, "initialize default SRS")
} }
@ -75,12 +75,12 @@ func (v *MemoryLoadBalancer) Initialize(ctx context.Context) error {
return nil return nil
} }
func (v *MemoryLoadBalancer) Update(ctx context.Context, server *OriginServer) error { func (v *memoryLoadBalancer) Update(ctx context.Context, server *OriginServer) error {
v.servers.Store(server.ID(), server) v.servers.Store(server.ID(), server)
return nil return nil
} }
func (v *MemoryLoadBalancer) Pick(ctx context.Context, streamURL string) (*OriginServer, error) { func (v *memoryLoadBalancer) Pick(ctx context.Context, streamURL string) (*OriginServer, error) {
// Always proxy to the same server for the same stream URL. // Always proxy to the same server for the same stream URL.
if server, ok := v.picked.Load(streamURL); ok { if server, ok := v.picked.Load(streamURL); ok {
return server, nil return server, nil
@ -115,7 +115,7 @@ func (v *MemoryLoadBalancer) Pick(ctx context.Context, streamURL string) (*Origi
return server, nil return server, nil
} }
func (v *MemoryLoadBalancer) LoadHLSBySPBHID(ctx context.Context, spbhid string) (HLSPlayStream, error) { func (v *memoryLoadBalancer) LoadHLSBySPBHID(ctx context.Context, spbhid string) (HLSPlayStream, error) {
// Load the HLS streaming for the SPBHID, for TS files. // Load the HLS streaming for the SPBHID, for TS files.
if actual, ok := v.hlsSPBHID.Load(spbhid); !ok { if actual, ok := v.hlsSPBHID.Load(spbhid); !ok {
return nil, errors.Errorf("no HLS streaming for SPBHID %v", spbhid) return nil, errors.Errorf("no HLS streaming for SPBHID %v", spbhid)
@ -124,7 +124,7 @@ func (v *MemoryLoadBalancer) LoadHLSBySPBHID(ctx context.Context, spbhid string)
} }
} }
func (v *MemoryLoadBalancer) LoadOrStoreHLS(ctx context.Context, streamURL string, value HLSPlayStream) (HLSPlayStream, error) { func (v *memoryLoadBalancer) LoadOrStoreHLS(ctx context.Context, streamURL string, value HLSPlayStream) (HLSPlayStream, error) {
// Update the HLS streaming for the stream URL, for M3u8. // Update the HLS streaming for the stream URL, for M3u8.
actual, _ := v.hlsStreamURL.LoadOrStore(streamURL, value) actual, _ := v.hlsStreamURL.LoadOrStore(streamURL, value)
if actual == nil { if actual == nil {
@ -137,7 +137,7 @@ func (v *MemoryLoadBalancer) LoadOrStoreHLS(ctx context.Context, streamURL strin
return actual, nil return actual, nil
} }
func (v *MemoryLoadBalancer) StoreWebRTC(ctx context.Context, streamURL string, value RTCConnection) error { func (v *memoryLoadBalancer) StoreWebRTC(ctx context.Context, streamURL string, value RTCConnection) error {
// Update the WebRTC streaming for the stream URL. // Update the WebRTC streaming for the stream URL.
v.rtcStreamURL.Store(streamURL, value) v.rtcStreamURL.Store(streamURL, value)
@ -146,7 +146,7 @@ func (v *MemoryLoadBalancer) StoreWebRTC(ctx context.Context, streamURL string,
return nil return nil
} }
func (v *MemoryLoadBalancer) LoadWebRTCByUfrag(ctx context.Context, ufrag string) (RTCConnection, error) { func (v *memoryLoadBalancer) LoadWebRTCByUfrag(ctx context.Context, ufrag string) (RTCConnection, error) {
if actual, ok := v.rtcUfrag.Load(ufrag); !ok { if actual, ok := v.rtcUfrag.Load(ufrag); !ok {
return nil, errors.Errorf("no WebRTC streaming for ufrag %v", ufrag) return nil, errors.Errorf("no WebRTC streaming for ufrag %v", ufrag)
} else { } else {

34
internal/lb/redis.go
View File

@ -19,8 +19,8 @@ import (
"srsx/internal/logger" "srsx/internal/logger"
) )
// RedisLoadBalancer stores state in Redis. // redisLoadBalancer stores state in Redis.
type RedisLoadBalancer struct { type redisLoadBalancer struct {
// The environment interface. // The environment interface.
environment env.ProxyEnvironment environment env.ProxyEnvironment
// The redis client sdk. // The redis client sdk.
@ -29,12 +29,12 @@ type RedisLoadBalancer struct {
// NewRedisLoadBalancer creates a new Redis-based load balancer. // NewRedisLoadBalancer creates a new Redis-based load balancer.
func NewRedisLoadBalancer(environment env.ProxyEnvironment) OriginLoadBalancer { func NewRedisLoadBalancer(environment env.ProxyEnvironment) OriginLoadBalancer {
return &RedisLoadBalancer{ return &redisLoadBalancer{
environment: environment, environment: environment,
} }
} }
func (v *RedisLoadBalancer) Initialize(ctx context.Context) error { func (v *redisLoadBalancer) Initialize(ctx context.Context) error {
redisDatabase, err := strconv.Atoi(v.environment.RedisDB()) redisDatabase, err := strconv.Atoi(v.environment.RedisDB())
if err != nil { if err != nil {
return errors.Wrapf(err, "invalid PROXY_REDIS_DB %v", v.environment.RedisDB()) return errors.Wrapf(err, "invalid PROXY_REDIS_DB %v", v.environment.RedisDB())
@ -52,7 +52,7 @@ func (v *RedisLoadBalancer) Initialize(ctx context.Context) error {
} }
logger.Debug(ctx, "RedisLB: connected to redis %v ok", rdb.String()) logger.Debug(ctx, "RedisLB: connected to redis %v ok", rdb.String())
server, err := NewDefaultSRSForDebugging(v.environment) server, err := NewDefaultOriginServerForDebugging(v.environment)
if err != nil { if err != nil {
return errors.Wrapf(err, "initialize default SRS") return errors.Wrapf(err, "initialize default SRS")
} }
@ -80,7 +80,7 @@ func (v *RedisLoadBalancer) Initialize(ctx context.Context) error {
return nil return nil
} }
func (v *RedisLoadBalancer) Update(ctx context.Context, server *OriginServer) error { func (v *redisLoadBalancer) Update(ctx context.Context, server *OriginServer) error {
b, err := json.Marshal(server) b, err := json.Marshal(server)
if err != nil { if err != nil {
return errors.Wrapf(err, "marshal server %+v", server) return errors.Wrapf(err, "marshal server %+v", server)
@ -130,7 +130,7 @@ func (v *RedisLoadBalancer) Update(ctx context.Context, server *OriginServer) er
return nil return nil
} }
func (v *RedisLoadBalancer) Pick(ctx context.Context, streamURL string) (*OriginServer, error) { func (v *redisLoadBalancer) Pick(ctx context.Context, streamURL string) (*OriginServer, error) {
key := fmt.Sprintf("srs-proxy-url:%v", streamURL) key := fmt.Sprintf("srs-proxy-url:%v", streamURL)
// Always proxy to the same server for the same stream URL. // Always proxy to the same server for the same stream URL.
@ -188,7 +188,7 @@ func (v *RedisLoadBalancer) Pick(ctx context.Context, streamURL string) (*Origin
return &server, nil return &server, nil
} }
func (v *RedisLoadBalancer) LoadHLSBySPBHID(ctx context.Context, spbhid string) (HLSPlayStream, error) { func (v *redisLoadBalancer) LoadHLSBySPBHID(ctx context.Context, spbhid string) (HLSPlayStream, error) {
key := v.redisKeySPBHID(spbhid) key := v.redisKeySPBHID(spbhid)
b, err := v.rdb.Get(ctx, key).Bytes() b, err := v.rdb.Get(ctx, key).Bytes()
@ -208,7 +208,7 @@ func (v *RedisLoadBalancer) LoadHLSBySPBHID(ctx context.Context, spbhid string)
return nil, errors.Errorf("Redis load balancer cannot deserialize interface types") return nil, errors.Errorf("Redis load balancer cannot deserialize interface types")
} }
func (v *RedisLoadBalancer) LoadOrStoreHLS(ctx context.Context, streamURL string, value HLSPlayStream) (HLSPlayStream, error) { func (v *redisLoadBalancer) LoadOrStoreHLS(ctx context.Context, streamURL string, value HLSPlayStream) (HLSPlayStream, error) {
b, err := json.Marshal(value) b, err := json.Marshal(value)
if err != nil { if err != nil {
return nil, errors.Wrapf(err, "marshal HLS %v", value) return nil, errors.Wrapf(err, "marshal HLS %v", value)
@ -229,7 +229,7 @@ func (v *RedisLoadBalancer) LoadOrStoreHLS(ctx context.Context, streamURL string
return value, nil return value, nil
} }
func (v *RedisLoadBalancer) StoreWebRTC(ctx context.Context, streamURL string, value RTCConnection) error { func (v *redisLoadBalancer) StoreWebRTC(ctx context.Context, streamURL string, value RTCConnection) error {
b, err := json.Marshal(value) b, err := json.Marshal(value)
if err != nil { if err != nil {
return errors.Wrapf(err, "marshal WebRTC %v", value) return errors.Wrapf(err, "marshal WebRTC %v", value)
@ -249,7 +249,7 @@ func (v *RedisLoadBalancer) StoreWebRTC(ctx context.Context, streamURL string, v
return nil return nil
} }
func (v *RedisLoadBalancer) LoadWebRTCByUfrag(ctx context.Context, ufrag string) (RTCConnection, error) { func (v *redisLoadBalancer) LoadWebRTCByUfrag(ctx context.Context, ufrag string) (RTCConnection, error) {
key := v.redisKeyUfrag(ufrag) key := v.redisKeyUfrag(ufrag)
b, err := v.rdb.Get(ctx, key).Bytes() b, err := v.rdb.Get(ctx, key).Bytes()
@ -267,26 +267,26 @@ func (v *RedisLoadBalancer) LoadWebRTCByUfrag(ctx context.Context, ufrag string)
return nil, errors.Errorf("Redis load balancer cannot deserialize interface types") return nil, errors.Errorf("Redis load balancer cannot deserialize interface types")
} }
func (v *RedisLoadBalancer) redisKeyUfrag(ufrag string) string { func (v *redisLoadBalancer) redisKeyUfrag(ufrag string) string {
return fmt.Sprintf("srs-proxy-ufrag:%v", ufrag) return fmt.Sprintf("srs-proxy-ufrag:%v", ufrag)
} }
func (v *RedisLoadBalancer) redisKeyRTC(streamURL string) string { func (v *redisLoadBalancer) redisKeyRTC(streamURL string) string {
return fmt.Sprintf("srs-proxy-rtc:%v", streamURL) return fmt.Sprintf("srs-proxy-rtc:%v", streamURL)
} }
func (v *RedisLoadBalancer) redisKeySPBHID(spbhid string) string { func (v *redisLoadBalancer) redisKeySPBHID(spbhid string) string {
return fmt.Sprintf("srs-proxy-spbhid:%v", spbhid) return fmt.Sprintf("srs-proxy-spbhid:%v", spbhid)
} }
func (v *RedisLoadBalancer) redisKeyHLS(streamURL string) string { func (v *redisLoadBalancer) redisKeyHLS(streamURL string) string {
return fmt.Sprintf("srs-proxy-hls:%v", streamURL) return fmt.Sprintf("srs-proxy-hls:%v", streamURL)
} }
func (v *RedisLoadBalancer) redisKeyServer(serverID string) string { func (v *redisLoadBalancer) redisKeyServer(serverID string) string {
return fmt.Sprintf("srs-proxy-server:%v", serverID) return fmt.Sprintf("srs-proxy-server:%v", serverID)
} }
func (v *RedisLoadBalancer) redisKeyServers() string { func (v *redisLoadBalancer) redisKeyServers() string {
return fmt.Sprintf("srs-proxy-all-servers") return fmt.Sprintf("srs-proxy-all-servers")
} }

7
internal/proxy/api.go
View File

@ -147,6 +147,8 @@ func (v *httpAPIProxyServer) Run(ctx context.Context) error {
type systemAPI struct { type systemAPI struct {
// The environment interface. // The environment interface.
environment env.ProxyEnvironment environment env.ProxyEnvironment
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
// The underlayer HTTP server. // The underlayer HTTP server.
server *http.Server server *http.Server
// The gracefully quit timeout, wait server to quit. // The gracefully quit timeout, wait server to quit.
@ -155,9 +157,10 @@ type systemAPI struct {
wg sync.WaitGroup wg sync.WaitGroup
} }
func NewSystemAPI(environment env.ProxyEnvironment, gracefulQuitTimeout time.Duration) *systemAPI { func NewSystemAPI(environment env.ProxyEnvironment, loadBalancer lb.OriginLoadBalancer, gracefulQuitTimeout time.Duration) *systemAPI {
v := &systemAPI{ v := &systemAPI{
environment: environment, environment: environment,
loadBalancer: loadBalancer,
gracefulQuitTimeout: gracefulQuitTimeout, gracefulQuitTimeout: gracefulQuitTimeout,
} }
return v return v
@ -262,7 +265,7 @@ func (v *systemAPI) Run(ctx context.Context) error {
srs.SRT, srs.RTC = srt, rtc srs.SRT, srs.RTC = srt, rtc
srs.UpdatedAt = time.Now() srs.UpdatedAt = time.Now()
}) })
if err := lb.SrsLoadBalancer.Update(ctx, server); err != nil { if err := v.loadBalancer.Update(ctx, server); err != nil {
return errors.Wrapf(err, "update SRS server %+v", server) return errors.Wrapf(err, "update SRS server %+v", server)
} }

19
internal/proxy/http.go
View File

@ -34,6 +34,8 @@ type HTTPStreamProxyServer interface {
type httpStreamProxyServer struct { type httpStreamProxyServer struct {
// The environment interface. // The environment interface.
environment env.ProxyEnvironment environment env.ProxyEnvironment
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
// The underlayer HTTP server. // The underlayer HTTP server.
server *http.Server server *http.Server
// The gracefully quit timeout, wait server to quit. // The gracefully quit timeout, wait server to quit.
@ -42,9 +44,10 @@ type httpStreamProxyServer struct {
wg stdSync.WaitGroup wg stdSync.WaitGroup
} }
func NewHTTPStreamProxyServer(environment env.ProxyEnvironment, gracefulQuitTimeout time.Duration) HTTPStreamProxyServer { func NewHTTPStreamProxyServer(environment env.ProxyEnvironment, loadBalancer lb.OriginLoadBalancer, gracefulQuitTimeout time.Duration) HTTPStreamProxyServer {
v := &httpStreamProxyServer{ v := &httpStreamProxyServer{
environment: environment, environment: environment,
loadBalancer: loadBalancer,
gracefulQuitTimeout: gracefulQuitTimeout, gracefulQuitTimeout: gracefulQuitTimeout,
} }
return v return v
@ -128,7 +131,8 @@ func (v *httpStreamProxyServer) Run(ctx context.Context) error {
return return
} }
stream, _ := lb.SrsLoadBalancer.LoadOrStoreHLS(ctx, streamURL, newHLSPlayStream(func(s *hlsPlayStream) { stream, _ := v.loadBalancer.LoadOrStoreHLS(ctx, streamURL, newHLSPlayStream(func(s *hlsPlayStream) {
s.loadBalancer = v.loadBalancer
s.SRSProxyBackendHLSID = logger.GenerateContextID() s.SRSProxyBackendHLSID = logger.GenerateContextID()
s.StreamURL, s.FullURL = streamURL, fullURL s.StreamURL, s.FullURL = streamURL, fullURL
})) }))
@ -142,7 +146,7 @@ func (v *httpStreamProxyServer) Run(ctx context.Context) error {
strings.HasSuffix(r.URL.Path, ".ts") { strings.HasSuffix(r.URL.Path, ".ts") {
// If SPBHID is specified, it must be a HLS stream client. // If SPBHID is specified, it must be a HLS stream client.
if srsProxyBackendID := r.URL.Query().Get("spbhid"); srsProxyBackendID != "" { if srsProxyBackendID := r.URL.Query().Get("spbhid"); srsProxyBackendID != "" {
if stream, err := lb.SrsLoadBalancer.LoadHLSBySPBHID(ctx, srsProxyBackendID); err != nil { if stream, err := v.loadBalancer.LoadHLSBySPBHID(ctx, srsProxyBackendID); err != nil {
http.Error(w, fmt.Sprintf("load stream by spbhid %v", srsProxyBackendID), http.StatusBadRequest) http.Error(w, fmt.Sprintf("load stream by spbhid %v", srsProxyBackendID), http.StatusBadRequest)
} else { } else {
stream.Initialize(ctx).(*hlsPlayStream).ServeHTTP(w, r) stream.Initialize(ctx).(*hlsPlayStream).ServeHTTP(w, r)
@ -153,6 +157,7 @@ func (v *httpStreamProxyServer) Run(ctx context.Context) error {
// Use HTTP pseudo streaming to proxy the request. // Use HTTP pseudo streaming to proxy the request.
newHTTPFlvTsConnection(func(c *httpFlvTsConnection) { newHTTPFlvTsConnection(func(c *httpFlvTsConnection) {
c.ctx = ctx c.ctx = ctx
c.loadBalancer = v.loadBalancer
}).ServeHTTP(w, r) }).ServeHTTP(w, r)
return return
} }
@ -196,6 +201,8 @@ func (v *httpStreamProxyServer) Run(ctx context.Context) error {
type httpFlvTsConnection struct { type httpFlvTsConnection struct {
// The context for HTTP streaming. // The context for HTTP streaming.
ctx context.Context ctx context.Context
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
} }
func newHTTPFlvTsConnection(opts ...func(*httpFlvTsConnection)) *httpFlvTsConnection { func newHTTPFlvTsConnection(opts ...func(*httpFlvTsConnection)) *httpFlvTsConnection {
@ -233,7 +240,7 @@ func (v *httpFlvTsConnection) serve(ctx context.Context, w http.ResponseWriter,
} }
// Pick a backend SRS server to proxy the RTMP stream. // Pick a backend SRS server to proxy the RTMP stream.
backend, err := lb.SrsLoadBalancer.Pick(ctx, streamURL) backend, err := v.loadBalancer.Pick(ctx, streamURL)
if err != nil { if err != nil {
return errors.Wrapf(err, "pick backend for %v", streamURL) return errors.Wrapf(err, "pick backend for %v", streamURL)
} }
@ -303,6 +310,8 @@ func (v *httpFlvTsConnection) serveByBackend(ctx context.Context, w http.Respons
type hlsPlayStream struct { type hlsPlayStream struct {
// The context for HLS streaming. // The context for HLS streaming.
ctx context.Context ctx context.Context
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
// The spbhid, used to identify the backend server. // The spbhid, used to identify the backend server.
SRSProxyBackendHLSID string `json:"spbhid"` SRSProxyBackendHLSID string `json:"spbhid"`
@ -351,7 +360,7 @@ func (v *hlsPlayStream) serve(ctx context.Context, w http.ResponseWriter, r *htt
} }
// Pick a backend SRS server to proxy the RTMP stream. // Pick a backend SRS server to proxy the RTMP stream.
backend, err := lb.SrsLoadBalancer.Pick(ctx, streamURL) backend, err := v.loadBalancer.Pick(ctx, streamURL)
if err != nil { if err != nil {
return errors.Wrapf(err, "pick backend for %v", streamURL) return errors.Wrapf(err, "pick backend for %v", streamURL)
} }

25
internal/proxy/rtc.go
View File

@ -36,6 +36,8 @@ type WebRTCProxyServer interface {
type webRTCProxyServer struct { type webRTCProxyServer struct {
// The environment interface. // The environment interface.
environment env.ProxyEnvironment environment env.ProxyEnvironment
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
// The UDP listener for WebRTC server. // The UDP listener for WebRTC server.
listener *net.UDPConn listener *net.UDPConn
@ -51,11 +53,12 @@ type webRTCProxyServer struct {
wg stdSync.WaitGroup wg stdSync.WaitGroup
} }
func NewWebRTCProxyServer(environment env.ProxyEnvironment, opts ...func(*webRTCProxyServer)) WebRTCProxyServer { func NewWebRTCProxyServer(environment env.ProxyEnvironment, loadBalancer lb.OriginLoadBalancer, opts ...func(*webRTCProxyServer)) WebRTCProxyServer {
v := &webRTCProxyServer{ v := &webRTCProxyServer{
environment: environment, environment: environment,
usernames: sync.NewMap[string, *rtcConnection](), loadBalancer: loadBalancer,
addresses: sync.NewMap[string, *rtcConnection](), usernames: sync.NewMap[string, *rtcConnection](),
addresses: sync.NewMap[string, *rtcConnection](),
} }
for _, opt := range opts { for _, opt := range opts {
opt(v) opt(v)
@ -97,7 +100,7 @@ func (v *webRTCProxyServer) HandleApiForWHIP(ctx context.Context, w http.Respons
} }
// Pick a backend SRS server to proxy the RTMP stream. // Pick a backend SRS server to proxy the RTMP stream.
backend, err := lb.SrsLoadBalancer.Pick(ctx, streamURL) backend, err := v.loadBalancer.Pick(ctx, streamURL)
if err != nil { if err != nil {
return errors.Wrapf(err, "pick backend for %v", streamURL) return errors.Wrapf(err, "pick backend for %v", streamURL)
} }
@ -134,7 +137,7 @@ func (v *webRTCProxyServer) HandleApiForWHEP(ctx context.Context, w http.Respons
} }
// Pick a backend SRS server to proxy the RTMP stream. // Pick a backend SRS server to proxy the RTMP stream.
backend, err := lb.SrsLoadBalancer.Pick(ctx, streamURL) backend, err := v.loadBalancer.Pick(ctx, streamURL)
if err != nil { if err != nil {
return errors.Wrapf(err, "pick backend for %v", streamURL) return errors.Wrapf(err, "pick backend for %v", streamURL)
} }
@ -226,7 +229,8 @@ func (v *webRTCProxyServer) proxyApiToBackend(
RemoteICEUfrag: remoteICEUfrag, RemoteICEPwd: remoteICEPwd, RemoteICEUfrag: remoteICEUfrag, RemoteICEPwd: remoteICEPwd,
LocalICEUfrag: localICEUfrag, LocalICEPwd: localICEPwd, LocalICEUfrag: localICEUfrag, LocalICEPwd: localICEPwd,
} }
if err := lb.SrsLoadBalancer.StoreWebRTC(ctx, streamURL, newRTCConnection(func(c *rtcConnection) { if err := v.loadBalancer.StoreWebRTC(ctx, streamURL, newRTCConnection(func(c *rtcConnection) {
c.loadBalancer = v.loadBalancer
c.StreamURL, c.Ufrag = streamURL, icePair.Ufrag() c.StreamURL, c.Ufrag = streamURL, icePair.Ufrag()
c.Initialize(ctx, v.listener) c.Initialize(ctx, v.listener)
@ -315,10 +319,11 @@ func (v *webRTCProxyServer) handleClientUDP(ctx context.Context, addr *net.UDPAd
} }
// Load connection by username. // Load connection by username.
if s, err := lb.SrsLoadBalancer.LoadWebRTCByUfrag(ctx, pkt.Username); err != nil { if s, err := v.loadBalancer.LoadWebRTCByUfrag(ctx, pkt.Username); err != nil {
return errors.Wrapf(err, "load webrtc by ufrag %v", pkt.Username) return errors.Wrapf(err, "load webrtc by ufrag %v", pkt.Username)
} else { } else {
connection = s.(*rtcConnection).Initialize(ctx, v.listener) connection = s.(*rtcConnection).Initialize(ctx, v.listener)
connection.loadBalancer = v.loadBalancer
logger.Debug(ctx, "Create WebRTC connection by ufrag=%v, stream=%v", pkt.Username, connection.StreamURL) logger.Debug(ctx, "Create WebRTC connection by ufrag=%v, stream=%v", pkt.Username, connection.StreamURL)
} }
@ -366,6 +371,8 @@ func (v *webRTCProxyServer) handleClientUDP(ctx context.Context, addr *net.UDPAd
type rtcConnection struct { type rtcConnection struct {
// The stream context for WebRTC streaming. // The stream context for WebRTC streaming.
ctx context.Context ctx context.Context
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
// The stream URL in vhost/app/stream schema. // The stream URL in vhost/app/stream schema.
StreamURL string `json:"stream_url"` StreamURL string `json:"stream_url"`
@ -450,7 +457,7 @@ func (v *rtcConnection) connectBackend(ctx context.Context) error {
} }
// Pick a backend SRS server to proxy the RTC stream. // Pick a backend SRS server to proxy the RTC stream.
backend, err := lb.SrsLoadBalancer.Pick(ctx, v.StreamURL) backend, err := v.loadBalancer.Pick(ctx, v.StreamURL)
if err != nil { if err != nil {
return errors.Wrapf(err, "pick backend") return errors.Wrapf(err, "pick backend")
} }

17
internal/proxy/rtmp.go
View File

@ -31,14 +31,16 @@ type RTMPProxyServer interface {
type rtmpProxyServer struct { type rtmpProxyServer struct {
// The environment interface. // The environment interface.
environment env.ProxyEnvironment environment env.ProxyEnvironment
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
// The TCP listener for RTMP server. // The TCP listener for RTMP server.
listener *net.TCPListener listener *net.TCPListener
// The wait group for all goroutines. // The wait group for all goroutines.
wg sync.WaitGroup wg sync.WaitGroup
} }
func NewRTMPProxyServer(environment env.ProxyEnvironment, opts ...func(*rtmpProxyServer)) RTMPProxyServer { func NewRTMPProxyServer(environment env.ProxyEnvironment, loadBalancer lb.OriginLoadBalancer, opts ...func(*rtmpProxyServer)) RTMPProxyServer {
v := &rtmpProxyServer{environment: environment} v := &rtmpProxyServer{environment: environment, loadBalancer: loadBalancer}
for _, opt := range opts { for _, opt := range opts {
opt(v) opt(v)
} }
@ -102,7 +104,9 @@ func (v *rtmpProxyServer) Run(ctx context.Context) error {
} }
} }
rc := newRTMPConnection() rc := newRTMPConnection(func(c *rtmpConnection) {
c.loadBalancer = v.loadBalancer
})
if err := rc.serve(ctx, conn); err != nil { if err := rc.serve(ctx, conn); err != nil {
handleErr(err) handleErr(err)
} else { } else {
@ -122,6 +126,8 @@ func (v *rtmpProxyServer) Run(ctx context.Context) error {
// then proxy to the corresponding backend server. All state is in the RTMP request, so this // then proxy to the corresponding backend server. All state is in the RTMP request, so this
// connection is stateless. // connection is stateless.
type rtmpConnection struct { type rtmpConnection struct {
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
} }
func newRTMPConnection(opts ...func(*rtmpConnection)) *rtmpConnection { func newRTMPConnection(opts ...func(*rtmpConnection)) *rtmpConnection {
@ -296,6 +302,7 @@ func (v *rtmpConnection) serve(ctx context.Context, conn *net.TCPConn) error {
// Find a backend SRS server to proxy the RTMP stream. // Find a backend SRS server to proxy the RTMP stream.
backend = newRTMPClientToBackend(func(client *rtmpClientToBackend) { backend = newRTMPClientToBackend(func(client *rtmpClientToBackend) {
client.typ = clientType client.typ = clientType
client.loadBalancer = v.loadBalancer
}) })
defer backend.Close() defer backend.Close()
@ -429,6 +436,8 @@ type rtmpClientToBackend struct {
client rtmp.Protocol client rtmp.Protocol
// The stream type. // The stream type.
typ RTMPClientType typ RTMPClientType
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
} }
func newRTMPClientToBackend(opts ...func(*rtmpClientToBackend)) *rtmpClientToBackend { func newRTMPClientToBackend(opts ...func(*rtmpClientToBackend)) *rtmpClientToBackend {
@ -454,7 +463,7 @@ func (v *rtmpClientToBackend) Connect(ctx context.Context, tcUrl, streamName str
} }
// Pick a backend SRS server to proxy the RTMP stream. // Pick a backend SRS server to proxy the RTMP stream.
backend, err := lb.SrsLoadBalancer.Pick(ctx, streamURL) backend, err := v.loadBalancer.Pick(ctx, streamURL)
if err != nil { if err != nil {
return errors.Wrapf(err, "pick backend for %v", streamURL) return errors.Wrapf(err, "pick backend for %v", streamURL)
} }

16
internal/proxy/srt.go
View File

@ -27,6 +27,8 @@ import (
type srsSRTProxyServer struct { type srsSRTProxyServer struct {
// The environment interface. // The environment interface.
environment env.ProxyEnvironment environment env.ProxyEnvironment
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
// The UDP listener for SRT server. // The UDP listener for SRT server.
listener *net.UDPConn listener *net.UDPConn
@ -39,11 +41,12 @@ type srsSRTProxyServer struct {
wg stdSync.WaitGroup wg stdSync.WaitGroup
} }
func NewSRSSRTProxyServer(environment env.ProxyEnvironment, opts ...func(*srsSRTProxyServer)) *srsSRTProxyServer { func NewSRSSRTProxyServer(environment env.ProxyEnvironment, loadBalancer lb.OriginLoadBalancer, opts ...func(*srsSRTProxyServer)) *srsSRTProxyServer {
v := &srsSRTProxyServer{ v := &srsSRTProxyServer{
environment: environment, environment: environment,
start: time.Now(), loadBalancer: loadBalancer,
sockets: sync.NewMap[uint32, *SRTConnection](), start: time.Now(),
sockets: sync.NewMap[uint32, *SRTConnection](),
} }
for _, opt := range opts { for _, opt := range opts {
@ -127,6 +130,7 @@ func (v *srsSRTProxyServer) handleClientUDP(ctx context.Context, addr *net.UDPAd
conn, ok := v.sockets.LoadOrStore(socketID, NewSRTConnection(func(c *SRTConnection) { conn, ok := v.sockets.LoadOrStore(socketID, NewSRTConnection(func(c *SRTConnection) {
c.ctx = logger.WithContext(ctx) c.ctx = logger.WithContext(ctx)
c.listenerUDP, c.socketID = v.listener, socketID c.listenerUDP, c.socketID = v.listener, socketID
c.loadBalancer = v.loadBalancer
c.start = v.start c.start = v.start
})) }))
@ -158,6 +162,8 @@ func (v *srsSRTProxyServer) handleClientUDP(ctx context.Context, addr *net.UDPAd
type SRTConnection struct { type SRTConnection struct {
// The stream context for SRT connection. // The stream context for SRT connection.
ctx context.Context ctx context.Context
// The load balancer for origin servers.
loadBalancer lb.OriginLoadBalancer
// The current socket ID. // The current socket ID.
socketID uint32 socketID uint32
@ -356,7 +362,7 @@ func (v *SRTConnection) connectBackend(ctx context.Context, streamID string) err
} }
// Pick a backend SRS server to proxy the SRT stream. // Pick a backend SRS server to proxy the SRT stream.
backend, err := lb.SrsLoadBalancer.Pick(ctx, streamURL) backend, err := v.loadBalancer.Pick(ctx, streamURL)
if err != nil { if err != nil {
return errors.Wrapf(err, "pick backend for %v", streamURL) return errors.Wrapf(err, "pick backend for %v", streamURL)
} }