Initial version for 2023 NHKG broadcast.
This commit is contained in:
279
main.go
Normal file
279
main.go
Normal file
@ -0,0 +1,279 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"container/list"
|
||||||
|
"flag"
|
||||||
|
"fmt"
|
||||||
|
"net"
|
||||||
|
"net/http"
|
||||||
|
"net/http/pprof"
|
||||||
|
"net/url"
|
||||||
|
"os"
|
||||||
|
"os/signal"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"syscall"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type UpstreamConfig map[string]*url.URL
|
||||||
|
|
||||||
|
func (config *UpstreamConfig) String() string {
|
||||||
|
return fmt.Sprintf("%+v", (*map[string]*url.URL)(config))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (config *UpstreamConfig) Set(value string) error {
|
||||||
|
parts := strings.SplitN(value, "=", 2)
|
||||||
|
if len(parts) < 2 {
|
||||||
|
return fmt.Errorf("format error: not in name=url format: %s", value)
|
||||||
|
}
|
||||||
|
name := parts[0]
|
||||||
|
if _, ok := (*config)[name]; ok {
|
||||||
|
return fmt.Errorf("config error: remote %s already specified", name)
|
||||||
|
}
|
||||||
|
u, err := url.Parse(parts[1])
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("format error: unable to parse url %s: %w", parts[1], err)
|
||||||
|
}
|
||||||
|
(*config)[name] = u
|
||||||
|
fmt.Printf("Added remote %s with url %s\n", name, u)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
var listen = flag.String("listen", "unix:/var/run/tsb/tsb.sock", "Listen address of the service")
|
||||||
|
|
||||||
|
var upstreams = make(UpstreamConfig)
|
||||||
|
|
||||||
|
type UpstreamContext struct {
|
||||||
|
viewerCond *sync.Cond
|
||||||
|
viewer int64
|
||||||
|
upstreamChan chan []byte
|
||||||
|
clientsChanList *list.List
|
||||||
|
newClientsChanList *list.List
|
||||||
|
newChanLock *sync.Mutex
|
||||||
|
}
|
||||||
|
|
||||||
|
var upstreamContexts = make(map[string]*UpstreamContext)
|
||||||
|
var upstreamContextLock = &sync.Mutex{}
|
||||||
|
|
||||||
|
func GetUpstreamContext(name string) *UpstreamContext {
|
||||||
|
upstreamContextLock.Lock()
|
||||||
|
defer upstreamContextLock.Unlock()
|
||||||
|
ctx, ok := upstreamContexts[name]
|
||||||
|
if !ok {
|
||||||
|
ctx = &UpstreamContext{
|
||||||
|
viewerCond: sync.NewCond(&sync.Mutex{}),
|
||||||
|
viewer: 0,
|
||||||
|
upstreamChan: make(chan []byte),
|
||||||
|
clientsChanList: list.New(),
|
||||||
|
newClientsChanList: list.New(),
|
||||||
|
newChanLock: &sync.Mutex{},
|
||||||
|
}
|
||||||
|
upstreamContexts[name] = ctx
|
||||||
|
go upstreamFiber(name)
|
||||||
|
go broadcastFiber(name)
|
||||||
|
}
|
||||||
|
return ctx
|
||||||
|
}
|
||||||
|
|
||||||
|
type ClientChannel struct {
|
||||||
|
Notify chan struct{}
|
||||||
|
Data chan []byte
|
||||||
|
ctx *UpstreamContext
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewClientChannel(name string) *ClientChannel {
|
||||||
|
ctx := GetUpstreamContext(name)
|
||||||
|
ch := &ClientChannel{
|
||||||
|
Notify: make(chan struct{}),
|
||||||
|
Data: make(chan []byte, 1024),
|
||||||
|
ctx: ctx,
|
||||||
|
}
|
||||||
|
ctx.newChanLock.Lock()
|
||||||
|
ctx.newClientsChanList.PushBack(ch)
|
||||||
|
ctx.newChanLock.Unlock()
|
||||||
|
if atomic.AddInt64(&ctx.viewer, 1) == 1 {
|
||||||
|
ctx.viewerCond.Broadcast()
|
||||||
|
}
|
||||||
|
fmt.Printf("Current viewers for stream %s: %d\n", name, ctx.viewer)
|
||||||
|
return ch
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close the notify channel, to the server will stop sending data and remove this entry.
|
||||||
|
func (ch *ClientChannel) Close() {
|
||||||
|
close(ch.Notify)
|
||||||
|
atomic.AddInt64(&ch.ctx.viewer, -1)
|
||||||
|
}
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
flag.Var(&upstreams, "upstream", "Upstream in name=URL format, may specify multiple times.")
|
||||||
|
}
|
||||||
|
|
||||||
|
func upstreamFiber(name string) {
|
||||||
|
ctx := GetUpstreamContext(name)
|
||||||
|
client := &http.Client{}
|
||||||
|
req := &http.Request{
|
||||||
|
Method: "GET",
|
||||||
|
URL: upstreams[name],
|
||||||
|
Close: true,
|
||||||
|
}
|
||||||
|
for {
|
||||||
|
ctx.viewerCond.L.Lock()
|
||||||
|
if ctx.viewer == 0 {
|
||||||
|
fmt.Printf("No active viewers for stream %s, idling around :)\n", name)
|
||||||
|
ctx.viewerCond.Wait()
|
||||||
|
}
|
||||||
|
ctx.viewerCond.L.Unlock()
|
||||||
|
fmt.Printf("Connecting to upstream for stream %s ...\n", name)
|
||||||
|
r := (func() *http.Response {
|
||||||
|
for {
|
||||||
|
r, err := client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
fmt.Fprintf(os.Stderr, "Unable to request %s for stream %s: %v\n", upstreams[name], name, err)
|
||||||
|
time.Sleep(5 * time.Second)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if r.StatusCode != 200 {
|
||||||
|
fmt.Fprintf(os.Stderr, "Request to %s for stream %s failed with status code %d\n", upstreams[name], name, err)
|
||||||
|
time.Sleep(5 * time.Second)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
return r
|
||||||
|
}
|
||||||
|
})()
|
||||||
|
body := r.Body
|
||||||
|
for {
|
||||||
|
buff := make([]byte, 65536)
|
||||||
|
if ctx.viewer == 0 {
|
||||||
|
fmt.Printf("No active viewers for %s, stopping reading from upstream.\n", name)
|
||||||
|
break
|
||||||
|
}
|
||||||
|
n, err := body.Read(buff)
|
||||||
|
if err != nil {
|
||||||
|
fmt.Fprintf(os.Stderr, "Upstream for stream %s request reached EOF\n", name)
|
||||||
|
break
|
||||||
|
}
|
||||||
|
ctx.upstreamChan <- buff[:n]
|
||||||
|
}
|
||||||
|
body.Close()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func broadcastFiber(name string) {
|
||||||
|
ctx := GetUpstreamContext(name)
|
||||||
|
for {
|
||||||
|
chunk := <-ctx.upstreamChan
|
||||||
|
// Send to all existing clients
|
||||||
|
e := ctx.clientsChanList.Front()
|
||||||
|
for e != nil {
|
||||||
|
ch := e.Value.(*ClientChannel)
|
||||||
|
select {
|
||||||
|
case <-ch.Notify:
|
||||||
|
// Closed client, remove this entry
|
||||||
|
close(ch.Data)
|
||||||
|
next := e.Next()
|
||||||
|
ctx.clientsChanList.Remove(e)
|
||||||
|
e = next
|
||||||
|
case ch.Data <- chunk:
|
||||||
|
e = e.Next()
|
||||||
|
default:
|
||||||
|
e = e.Next()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Try to serve all new clients, but don't get locked up here :)
|
||||||
|
if ctx.newClientsChanList.Len() > 0 && ctx.newChanLock.TryLock() {
|
||||||
|
ctx.clientsChanList.PushBackList(ctx.newClientsChanList)
|
||||||
|
ctx.newClientsChanList.Init()
|
||||||
|
ctx.newChanLock.Unlock()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func clientHandler(w http.ResponseWriter, r *http.Request) {
|
||||||
|
remote := r.RemoteAddr
|
||||||
|
if r.Header.Get("X-Remote-Addr") != "" {
|
||||||
|
remote = r.Header.Get("X-Remote-Addr")
|
||||||
|
}
|
||||||
|
fmt.Printf("Client connection from %s accepted.\n", remote)
|
||||||
|
w.Header().Add("Content-Type", "video/MP2T")
|
||||||
|
w.WriteHeader(200)
|
||||||
|
name := r.URL.Path
|
||||||
|
if _, ok := upstreams[name]; !ok {
|
||||||
|
http.NotFound(w, r)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ch := NewClientChannel(name)
|
||||||
|
defer ch.Close()
|
||||||
|
for {
|
||||||
|
chunk := <-ch.Data
|
||||||
|
_, err := w.Write(chunk)
|
||||||
|
if err != nil {
|
||||||
|
fmt.Printf("Write to remote client %s failed, closing...\n", remote)
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func viewersHandler(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Add("Content-Type", "text/plain")
|
||||||
|
w.WriteHeader(200)
|
||||||
|
name := r.URL.Path
|
||||||
|
if name == "" {
|
||||||
|
totalViewer := int64(0)
|
||||||
|
upstreamContextLock.Lock()
|
||||||
|
for name, ctx := range upstreamContexts {
|
||||||
|
totalViewer += ctx.viewer
|
||||||
|
fmt.Fprintf(w, "%s: %d\n", name, ctx.viewer)
|
||||||
|
}
|
||||||
|
upstreamContextLock.Unlock()
|
||||||
|
fmt.Fprintf(w, "sum: %d\n", totalViewer)
|
||||||
|
} else {
|
||||||
|
if _, ok := upstreams[name]; !ok {
|
||||||
|
http.NotFound(w, r)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
upstreamContextLock.Lock()
|
||||||
|
ctx, ok := upstreamContexts[name]
|
||||||
|
if !ok {
|
||||||
|
fmt.Fprintf(w, "%s: 0\n", name)
|
||||||
|
} else {
|
||||||
|
fmt.Fprintf(w, "%s: %d\n", name, ctx.viewer)
|
||||||
|
}
|
||||||
|
upstreamContextLock.Unlock()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
flag.Parse()
|
||||||
|
var listener net.Listener
|
||||||
|
var err error
|
||||||
|
if strings.HasPrefix(*listen, "unix:") {
|
||||||
|
if listener, err = net.Listen("unix", (*listen)[5:]); err != nil {
|
||||||
|
panic(err)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
if listener, err = net.Listen("tcp", *listen); err != nil {
|
||||||
|
panic(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
handler := http.NewServeMux()
|
||||||
|
handler.Handle("/stream/", http.StripPrefix("/stream/", http.HandlerFunc(clientHandler)))
|
||||||
|
handler.Handle("/viewers/", http.StripPrefix("/viewers/", http.HandlerFunc(viewersHandler)))
|
||||||
|
handler.HandleFunc("/debug/pprof/", pprof.Index)
|
||||||
|
handler.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
|
||||||
|
handler.HandleFunc("/debug/pprof/profile", pprof.Profile)
|
||||||
|
handler.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
|
||||||
|
handler.HandleFunc("/debug/pprof/trace", pprof.Trace)
|
||||||
|
server := &http.Server{
|
||||||
|
Handler: handler,
|
||||||
|
}
|
||||||
|
go server.Serve(listener)
|
||||||
|
ch := make(chan os.Signal)
|
||||||
|
signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
|
||||||
|
select {
|
||||||
|
case sig := <-ch:
|
||||||
|
fmt.Fprintf(os.Stderr, "Caught signal %v, exiting...\n", sig)
|
||||||
|
os.Exit(0)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user