82 lines
1.5 KiB
Go
82 lines
1.5 KiB
Go
package services
|
|
|
|
import (
|
|
"bufio"
|
|
"io"
|
|
"log"
|
|
"sync"
|
|
)
|
|
|
|
type LogStreamer struct {
|
|
mu sync.RWMutex
|
|
subs map[string]map[string]chan string
|
|
}
|
|
|
|
func NewLogStreamer() *LogStreamer {
|
|
return &LogStreamer{
|
|
subs: make(map[string]map[string]chan string),
|
|
}
|
|
}
|
|
|
|
func (ls *LogStreamer) Subscribe(serverID, clientID string) chan string {
|
|
ls.mu.Lock()
|
|
defer ls.mu.Unlock()
|
|
|
|
if ls.subs[serverID] == nil {
|
|
ls.subs[serverID] = make(map[string]chan string)
|
|
}
|
|
|
|
ch := make(chan string, 256)
|
|
ls.subs[serverID][clientID] = ch
|
|
return ch
|
|
}
|
|
|
|
func (ls *LogStreamer) Unsubscribe(serverID, clientID string) {
|
|
ls.mu.Lock()
|
|
defer ls.mu.Unlock()
|
|
|
|
if ls.subs[serverID] != nil {
|
|
if ch, ok := ls.subs[serverID][clientID]; ok {
|
|
close(ch)
|
|
delete(ls.subs[serverID], clientID)
|
|
}
|
|
if len(ls.subs[serverID]) == 0 {
|
|
delete(ls.subs, serverID)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (ls *LogStreamer) Stream(serverID string, reader io.Reader, closeMsg string) {
|
|
scanner := bufio.NewScanner(reader)
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
ls.broadcast(serverID, line)
|
|
}
|
|
if err := scanner.Err(); err != nil {
|
|
log.Printf("log stream error for server %s: %v", serverID, err)
|
|
}
|
|
if closeMsg != "" {
|
|
ls.broadcast(serverID, closeMsg)
|
|
}
|
|
}
|
|
|
|
func (ls *LogStreamer) Broadcast(serverID, line string) {
|
|
ls.broadcast(serverID, line)
|
|
}
|
|
|
|
func (ls *LogStreamer) broadcast(serverID, line string) {
|
|
ls.mu.RLock()
|
|
defer ls.mu.RUnlock()
|
|
|
|
if ls.subs[serverID] == nil {
|
|
return
|
|
}
|
|
|
|
for _, ch := range ls.subs[serverID] {
|
|
select {
|
|
case ch <- line:
|
|
default:
|
|
}
|
|
}
|
|
}
|