restream-tools/player/main.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"
)

func getIp(r *http.Request) string {
	// Get the remote address
	var ip = r.RemoteAddr

	// If behind a proxy, get the real IP
	if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
		ip = fwd
	}

	return ip
}

type Viewer struct {
	LastActive   int64
	ActiveMirror int
}

type ViewCountRequest struct {
	ActiveMirror int `json:"active_mirror"`
}

type ViewCountResponse struct {
	TotalCount int         `json:"total_count"`
	ByMirror   map[int]int `json:"by_mirror"`
}

func main() {
	var addr = flag.String("addr", ":8080", "http service address")
	flag.Parse()

	var viewers = make(map[string]Viewer)
	var mirrorActivity = make(map[int]int)
	var lockViewers sync.Mutex
	var lockMirrorActivity sync.Mutex

	// Handle view counting
	http.HandleFunc("/viewcount", func(w http.ResponseWriter, r *http.Request) {
		// Get the current time
		var now = time.Now().Unix()

		// Get the remote address
		var ip = getIp(r)

		// Decode body JSON as a ViewerMetric
		var metric ViewCountRequest
		if err := json.NewDecoder(r.Body).Decode(&metric); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		// Update the last active time
		lockViewers.Lock()
		viewers[ip] = Viewer{
			LastActive:   now,
			ActiveMirror: metric.ActiveMirror,
		}
		var viewersCount = len(viewers)
		lockViewers.Unlock()

		// Respond
		lockMirrorActivity.Lock()
		var response = ViewCountResponse{
			TotalCount: viewersCount,
			ByMirror:   mirrorActivity,
		}
		lockMirrorActivity.Unlock()

		if err := json.NewEncoder(w).Encode(response); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	})

	// Periodically remove inactive viewers
	go func() {
		for {
			time.Sleep(time.Second)

			// Get the current time
			var now = time.Now().Unix()

			// Remove any viewers that have not been active in the last 10 seconds
			lockViewers.Lock()
			for ip, viewer := range viewers {
				if now-viewer.LastActive > 10 {
					delete(viewers, ip)
				}
			}
			lockViewers.Unlock()
		}
	}()

	// Periodically log the number of viewers
	go func() {
		// Prepare file to log to
		var f, ferr = os.OpenFile("viewers.log", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
		if ferr != nil {
			fmt.Println("Error opening file:", ferr)
			fmt.Println("viewer count will not be logged")
			return
		}

		// Loop forever
		var last_count = 0
		for {
			var now = time.Now().Format(time.RFC3339)
			lockViewers.Lock()
			var count = len(viewers)
			lockViewers.Unlock()

			// If the viewer count has changed, log it
			if count != last_count {
				fmt.Fprintf(f, "%s\t%d\n", now, count)
				last_count = count
			}

			time.Sleep(time.Second)
		}
	}()

	// Periodically aggregate viewer count per mirror
	go func() {
		// Prepare file to log to
		var f, ferr = os.OpenFile("viewers_per_mirror.log", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
		if ferr != nil {
			fmt.Println("Error opening file:", ferr)
			fmt.Println("mirror statistics will not be logged")
			return
		}

		// Loop forever
		for {
			var now = time.Now().Format(time.RFC3339)
			var counts = make(map[int]int)

			// Loop through all viewers
			lockViewers.Lock()
			for _, viewer := range viewers {
				// Increment the mirror count
				counts[viewer.ActiveMirror]++
			}
			lockViewers.Unlock()

			// Check if the counts have changed
			var changed = false
			lockMirrorActivity.Lock()
			for mirror, count := range counts {
				if count != mirrorActivity[mirror] {
					changed = true
					break
				}
			}
			lockMirrorActivity.Unlock()

			// If the counts have changed, log them
			if changed {
				for mirror, count := range counts {
					fmt.Fprintf(f, "%s\t%d\t%d\n", now, mirror, count)
				}
			}

			// Update the mirror activity
			lockMirrorActivity.Lock()
			mirrorActivity = counts
			lockMirrorActivity.Unlock()
			time.Sleep(time.Second)
		}
	}()

	// Handle static files
	http.Handle("/", http.FileServer(http.Dir("./static")))

	// Start the server
	fmt.Println("Starting server on", *addr)
	err := http.ListenAndServe(*addr, nil)
	if err != nil {
		fmt.Println(err)
	}
}