viewercount/main.go

98 lines
2.0 KiB
Go
Raw Normal View History

2020-12-28 01:26:55 +01:00
package main
import (
"bufio"
2020-12-28 02:22:33 +01:00
"context"
2020-12-28 01:26:55 +01:00
"fmt"
2020-12-28 02:22:33 +01:00
"io"
2020-12-28 01:26:55 +01:00
"log"
2020-12-28 02:22:33 +01:00
"net/http"
2020-12-28 01:26:55 +01:00
"os"
"regexp"
"strconv"
2020-12-28 02:22:33 +01:00
"github.com/docker/docker/api/types"
"github.com/docker/docker/client"
2020-12-28 01:26:55 +01:00
)
const COUNTER_SLOTS = 20
func main() {
2020-12-28 02:22:33 +01:00
containerName, ok := os.LookupEnv("CONTAINER_NAME")
if !ok {
log.Fatalf("unable to read CONTAINER_NAME")
}
2020-12-28 02:39:08 +01:00
listen, ok := os.LookupEnv("LISTEN")
if !ok {
listen = ":8080"
}
2020-12-28 02:22:33 +01:00
defaultHeaders := map[string]string{"User-Agent": "engine-api-cli-1.0"}
2020-12-28 02:39:33 +01:00
cli, err := client.NewClient("unix:///var/run/docker.sock", "v1.40", nil, defaultHeaders)
2020-12-28 02:22:33 +01:00
if err != nil {
log.Fatalf("unread: %s", err)
}
logs, err := cli.ContainerLogs(context.TODO(), containerName, types.ContainerLogsOptions{})
if err != nil {
log.Fatalf("unread: %s", err)
}
defer logs.Close()
2020-12-28 01:26:55 +01:00
counter := NewCounter()
2020-12-28 02:22:33 +01:00
go counter.Scan(logs)
2020-12-28 01:26:55 +01:00
http.HandleFunc("/", counter.MetricsHandler)
2020-12-28 02:39:08 +01:00
err = http.ListenAndServe(listen, nil)
2020-12-28 01:26:55 +01:00
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() {
2020-12-28 02:22:33 +01:00
c.Scan(os.Stdin)
}
func (c *Counter) Scan(in io.Reader) {
2020-12-28 01:26:55 +01:00
matcher := regexp.MustCompile(`/hls/stream-([0-9]*).ts`)
2020-12-28 02:22:33 +01:00
scanner := bufio.NewScanner(in)
2020-12-28 01:26:55 +01:00
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
}