56 lines
1.1 KiB
Go
56 lines
1.1 KiB
Go
package udp
|
|
|
|
import (
|
|
"net"
|
|
"sync"
|
|
|
|
mapset "github.com/deckarep/golang-set/v2"
|
|
)
|
|
|
|
type RoutingTable struct {
|
|
mu sync.RWMutex
|
|
routes map[string]mapset.Set[*net.UDPAddr]
|
|
}
|
|
|
|
func NewRoutingTable() *RoutingTable {
|
|
return &RoutingTable{
|
|
routes: make(map[string]mapset.Set[*net.UDPAddr]),
|
|
}
|
|
}
|
|
|
|
func (rt *RoutingTable) AddClient(channelID string, addr *net.UDPAddr) {
|
|
rt.mu.Lock()
|
|
defer rt.mu.Unlock()
|
|
|
|
if rt.routes[channelID] == nil {
|
|
rt.routes[channelID] = mapset.NewSet[*net.UDPAddr]()
|
|
}
|
|
rt.routes[channelID].Add(addr)
|
|
}
|
|
|
|
func (rt *RoutingTable) RemoveClient(channelID string, addr *net.UDPAddr) {
|
|
rt.mu.Lock()
|
|
defer rt.mu.Unlock()
|
|
|
|
if clients, exists := rt.routes[channelID]; exists {
|
|
clients.Remove(addr)
|
|
if clients.Cardinality() == 0 {
|
|
delete(rt.routes, channelID)
|
|
}
|
|
}
|
|
}
|
|
|
|
// GetClients returns the clients connected to the given channelID
|
|
// don't modify the returned set!
|
|
func (rt *RoutingTable) GetClients(channelID string) mapset.Set[*net.UDPAddr] {
|
|
rt.mu.RLock()
|
|
defer rt.mu.RUnlock()
|
|
|
|
clients, exists := rt.routes[channelID]
|
|
if !exists {
|
|
return nil
|
|
}
|
|
|
|
return clients
|
|
}
|