viewercount/main.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
}