viewercount/main.go

100 lines
2.0 KiB
Go

package main
import (
"bufio"
"context"
"fmt"
"io"
"log"
"net/http"
"os"
"regexp"
"strconv"
"github.com/docker/docker/api/types"
"github.com/docker/docker/client"
)
const COUNTER_SLOTS = 20
func main() {
containerName, ok := os.LookupEnv("CONTAINER_NAME")
if !ok {
log.Fatalf("unable to read CONTAINER_NAME")
}
listen, ok := os.LookupEnv("LISTEN")
if !ok {
listen = ":8080"
}
defaultHeaders := map[string]string{"User-Agent": "engine-api-cli-1.0"}
cli, err := client.NewClient("unix:///var/run/docker.sock", "v1.40", nil, defaultHeaders)
if err != nil {
log.Fatalf("unread: %s", err)
}
logs, err := cli.ContainerLogs(context.TODO(), containerName, types.ContainerLogsOptions{
ShowStdout: true,
})
if err != nil {
log.Fatalf("unread: %s", err)
}
defer logs.Close()
counter := NewCounter()
go counter.Scan(logs)
http.HandleFunc("/", counter.MetricsHandler)
err = http.ListenAndServe(listen, nil)
log.Fatalf("unable to listen: %s", err)
}
type Counter struct {
counters []int
names []string
current int
}
func NewCounter() *Counter {
return &Counter{
counters: make([]int, COUNTER_SLOTS),
names: make([]string, COUNTER_SLOTS),
}
}
func (c *Counter) MetricsHandler(w http.ResponseWriter, r *http.Request) {
w.Header().Add("Content-Type", "text/plain")
for i, name := range c.names {
fmt.Fprintf(w, "rc3_stream_count{name=%q} %d\n", name, c.counters[i])
}
fmt.Fprintf(w, "rc3_stream_current{name=\"%d\"} 1\n", c.current)
}
func (c *Counter) ScanStdin() {
c.Scan(os.Stdin)
}
func (c *Counter) Scan(in io.Reader) {
matcher := regexp.MustCompile(`/hls/stream-([0-9]*).ts`)
scanner := bufio.NewScanner(in)
scanner.Split(bufio.ScanLines)
for scanner.Scan() {
matches := matcher.FindStringSubmatch(scanner.Text())
if len(matches) == 2 {
c.countViewers(matches[1])
}
}
}
func (c *Counter) countViewers(part string) {
i, _ := strconv.Atoi(part)
mod := i % COUNTER_SLOTS
if c.names[mod] != part {
c.counters[mod] = 0
}
c.current = i
c.counters[mod] += 1
c.names[mod] = part
}