Create publisher for remote rtmp stream

rtmp/channel.go

@@ -0,0 +1,164 @@
+package rtmp
+
+import (
+	"context"
+	"net"
+	"net/url"
+	"path/filepath"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/datarhei/core/v16/session"
+	"github.com/datarhei/joy4/av"
+	"github.com/datarhei/joy4/av/pubsub"
+	"github.com/datarhei/joy4/format/rtmp"
+)
+
+type client struct {
+	conn      connection
+	id        string
+	createdAt time.Time
+
+	txbytes uint64
+	rxbytes uint64
+
+	collector session.Collector
+
+	cancel context.CancelFunc
+}
+
+func newClient(conn connection, id string, collector session.Collector) *client {
+	c := &client{
+		conn:      conn,
+		id:        id,
+		createdAt: time.Now(),
+
+		collector: collector,
+	}
+
+	var ctx context.Context
+	ctx, c.cancel = context.WithCancel(context.Background())
+
+	go c.ticker(ctx)
+
+	return c
+}
+
+func (c *client) ticker(ctx context.Context) {
+	ticker := time.NewTicker(1 * time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			txbytes := c.conn.TxBytes()
+			rxbytes := c.conn.RxBytes()
+
+			c.collector.Ingress(c.id, int64(rxbytes-c.rxbytes))
+			c.collector.Egress(c.id, int64(txbytes-c.txbytes))
+
+			c.txbytes = txbytes
+			c.rxbytes = rxbytes
+		}
+	}
+}
+
+func (c *client) Close() {
+	c.cancel()
+	c.conn.Close()
+}
+
+// channel represents a stream that is sent to the server
+type channel struct {
+	// The packet queue for the stream
+	queue *pubsub.Queue
+
+	// The metadata of the stream
+	streams []av.CodecData
+
+	// Whether the stream has an audio track
+	hasAudio bool
+
+	// Whether the stream has a video track
+	hasVideo bool
+
+	collector session.Collector
+	path      string
+
+	publisher  *client
+	subscriber map[string]*client
+	lock       sync.RWMutex
+
+	isProxy bool
+}
+
+func newChannel(conn connection, u *url.URL, remote net.Addr, streams []av.CodecData, isProxy bool, collector session.Collector) *channel {
+	ch := &channel{
+		path:       u.Path,
+		publisher:  newClient(conn, u.Path, collector),
+		subscriber: make(map[string]*client),
+		collector:  collector,
+		streams:    streams,
+		queue:      pubsub.NewQueue(),
+		isProxy:    isProxy,
+	}
+
+	addr := remote.String()
+	ip, _, _ := net.SplitHostPort(addr)
+
+	if collector.IsCollectableIP(ip) {
+		reference := strings.TrimSuffix(filepath.Base(u.Path), filepath.Ext(u.Path))
+		collector.RegisterAndActivate(u.Path, reference, "publish:"+u.Path, addr)
+	}
+
+	return ch
+}
+
+func (ch *channel) Close() {
+	if ch.publisher == nil {
+		return
+	}
+
+	ch.publisher.Close()
+	ch.publisher = nil
+
+	ch.queue.Close()
+}
+
+func (ch *channel) AddSubscriber(conn *rtmp.Conn) string {
+	addr := conn.NetConn().RemoteAddr().String()
+	ip, _, _ := net.SplitHostPort(addr)
+
+	client := newClient(conn, addr, ch.collector)
+
+	if ch.collector.IsCollectableIP(ip) {
+		reference := strings.TrimSuffix(filepath.Base(conn.URL.Path), filepath.Ext(conn.URL.Path))
+		ch.collector.RegisterAndActivate(addr, reference, "play:"+conn.URL.Path, addr)
+	}
+
+	ch.lock.Lock()
+	ch.subscriber[addr] = client
+	ch.lock.Unlock()
+
+	return addr
+}
+
+func (ch *channel) RemoveSubscriber(id string) {
+	ch.lock.Lock()
+	defer ch.lock.Unlock()
+
+	client := ch.subscriber[id]
+	if client != nil {
+		delete(ch.subscriber, id)
+		client.Close()
+	}
+
+	// If this is a proxied channel and the last subscriber leaves,
+	// close the channel.
+	if len(ch.subscriber) == 0 && ch.isProxy {
+		ch.Close()
+	}
+}