Added a websocket distributor

.gitignore

@@ -6,5 +6,6 @@ prompt_bak/
 /assert.debug
 /debug.doom
 /.env
+/vendor
 .zig-cache/
 zig-out/

pkg/v2/distributor/README.md

@@ -0,0 +1,11 @@
+# Websocket distributor
+
+This is completely dumb and just forwards messages to all downstreams.
+
+Auth is done by the downstreams which will disconnect a socket if it fails to
+auth, which will cause this to disconnect the upstream.
+
+## Usage
+
+Pass each `server:port` downstream on the command line.
+

pkg/v2/distributor/cmd/run/main.go

@@ -0,0 +1,39 @@
+package main
+
+import (
+	"flag"
+	"github.com/theprimeagen/vim-with-me/pkg/v2/distributor"
+	"os"
+	"strconv"
+	"log/slog"
+)
+
+func main() {
+	var port int
+	flag.IntVar(&port, "port", 0, "port to listen on")
+	flag.Parse()
+	downstreams := os.Args[1:]
+
+	if port == 0 {
+		portStr := os.Getenv("PORT")
+		if portStr == "" {
+			slog.Error("No port specified!")
+			os.Exit(1)
+		}
+		var err error
+		port, err = strconv.Atoi(portStr)
+		if err != nil {
+			slog.Error("Error converting port to int", "port", portStr, "err", err)
+			os.Exit(1)
+		}
+	}
+
+	authId := os.Getenv("AUTH_ID")
+	if authId == "" {
+		slog.Error("No auth id specified!")
+		os.Exit(1)
+	}
+
+	d := distributor.NewDistributor(port, authId, downstreams)
+	d.Run()
+}

pkg/v2/distributor/distributor.go

@@ -0,0 +1,115 @@
+package distributor
+
+import (
+	"fmt"
+	"github.com/gorilla/websocket"
+	"github.com/theprimeagen/vim-with-me/pkg/v2/assert"
+	"github.com/theprimeagen/vim-with-me/pkg/v2/metrics"
+	"log"
+	"log/slog"
+	"net/http"
+	"sync"
+)
+
+const (
+	upstreamConnectedMetricName       = "distributor_upstream_connected"
+	upstreamBytesReceivedMetricName   = "distributor_upstream_bytes_received"
+	upstreamBytesSentMetricName       = "distributor_upstream_bytes_sent"
+)
+
+type Distributor struct {
+	sync.Mutex
+
+	listenPort         int
+	authId             string
+	upstreamConnection *websocket.Conn
+	downstreams        []string
+	downstreamConns    []*Downstream
+	msgChan            chan []byte
+	stats              *metrics.Metrics
+}
+
+func NewDistributor(listenPort int, authId string, downstreams []string) *Distributor {
+	stats := metrics.New()
+	stats.Set(upstreamConnectedMetricName, 0)
+	stats.Set(upstreamBytesReceivedMetricName, 0)
+	stats.Set(upstreamBytesSentMetricName, 0)
+
+	return &Distributor{
+		listenPort:  listenPort,
+		authId:      authId,
+		downstreams: downstreams,
+		stats:       stats,
+	}
+}
+
+func (d *Distributor) Run() {
+	for _, addr := range d.downstreams {
+		d.downstreamConns = append(d.downstreamConns, NewDownstream(addr, d.stats))
+	}
+
+	http.HandleFunc("/metrics", func(w http.ResponseWriter, _ *http.Request) {
+		d.stats.WritePrometheus(w)
+	})
+
+	http.HandleFunc("/ws", d.handleIncomingConnection)
+
+	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Add("Location", "https://vim-with.me/")
+		w.WriteHeader(http.StatusTemporaryRedirect)
+		_, _ = w.Write([]byte("Not here, <a href=\"https://vim-with.me/\">over here</a>!"))
+	})
+
+	addr := fmt.Sprintf("0.0.0.0:%d", d.listenPort)
+	slog.Warn("listening and serving http", "http", addr)
+	err := http.ListenAndServe(addr, nil)
+
+	log.Fatal(err)
+}
+
+func (d *Distributor) handleIncomingConnection(w http.ResponseWriter, r *http.Request) {
+	func() {
+		d.Lock()
+		defer d.Unlock()
+
+		if d.upstreamConnection != nil {
+			// One connection at a time!
+			slog.Warn("Rejected connection, already have one",
+				"remote", r.RemoteAddr)
+			w.WriteHeader(http.StatusTooManyRequests)
+			return
+		}
+
+		slog.Info("New upstream connection", "remote", r.RemoteAddr)
+
+		upgrader := websocket.Upgrader{}
+		c, err := upgrader.Upgrade(w, r, nil)
+		assert.NoError(err, "unable to upgrade connection")
+
+		d.upstreamConnection = c
+
+		d.stats.Set(upstreamConnectedMetricName, 1)
+	}()
+
+	for {
+		mt, msg, err := d.upstreamConnection.ReadMessage()
+		if mt != websocket.BinaryMessage {
+			slog.Error("Upstream sent non-binary message, disconnecting")
+			break
+		}
+
+		if err != nil {
+			slog.Error("Upstream error, disconnecting", "err", err)
+			break
+		}
+
+		d.stats.Add(upstreamBytesReceivedMetricName, len(msg))
+		for _, downstream := range d.downstreamConns {
+			downstream.SendMessage(mt, msg)
+		}
+	}
+
+	d.stats.Set(upstreamConnectedMetricName, 0)
+	_ = d.upstreamConnection.Close()
+	d.upstreamConnection = nil
+}

pkg/v2/distributor/downstream.go

@@ -0,0 +1,75 @@
+package distributor
+
+import (
+	"github.com/gorilla/websocket"
+	"github.com/theprimeagen/vim-with-me/pkg/v2/assert"
+	"github.com/theprimeagen/vim-with-me/pkg/v2/metrics"
+	"log/slog"
+	"net/url"
+	"fmt"
+)
+
+const (
+	downstreamBytesReceivedMetricName = "distributor_downstream_bytes_received"
+	downstreamBytesSentMetricName     = "distributor_downstream_bytes_sent"
+)
+type Downstream struct {
+	addr  string
+	conn  *websocket.Conn
+	stats *metrics.Metrics
+}
+
+func NewDownstream(addr string, stats *metrics.Metrics) *Downstream {
+	ds := &Downstream{
+		addr:  addr,
+		stats: stats,
+	}
+	stats.Set(ds.makeMetricName(downstreamBytesReceivedMetricName), 0)
+	stats.Set(ds.makeMetricName(downstreamBytesSentMetricName), 0)
+	go ds.Run()
+	return ds
+}
+
+func (ds *Downstream) Run() {
+	var err error
+	for {
+		slog.Info("Connecting to downstream server", "addr", ds.addr)
+		u := url.URL{Scheme: "ws", Host: ds.addr, Path: "/ws"}
+		ds.conn, _, err = websocket.DefaultDialer.Dial(u.String(), nil)
+		assert.NoError(err, "unable to connect to downstream server "+ds.addr)
+
+		for {
+			// Discard any incoming messages
+			mt, msg, err := ds.conn.ReadMessage()
+			if err != nil || mt != websocket.BinaryMessage {
+				// Reconnect
+				slog.Warn("Downstream connection closed", "addr", ds.addr, "err", err)
+				break
+			}
+
+			ds.stats.Add(ds.makeMetricName(downstreamBytesReceivedMetricName), len(msg))
+		}
+
+		slog.Error("Downstream connection closed, reconnecting", "addr", ds.addr)
+
+		_ = ds.conn.Close()
+		ds.conn = nil
+	}
+}
+
+func (ds *Downstream) SendMessage(msgType int, msg []byte) {
+	err := ds.conn.WriteMessage(msgType, msg)
+	if err != nil {
+		// Reconnect
+		slog.Warn("Failed to send to downstream, closing", "addr", ds.addr, "err", err)
+		_ = ds.conn.Close()
+		ds.conn = nil
+		return
+	}
+
+	ds.stats.Add(ds.makeMetricName(downstreamBytesSentMetricName), len(msg))
+}
+
+func (ds *Downstream) makeMetricName(name string) string {
+	return fmt.Sprintf("%s{addr=\"%s\"}", name, ds.addr)
+}

pkg/v2/metrics/metrics.go

@@ -54,9 +54,13 @@ func (m *Metrics) SetIfGreater(key string, value int) {
 }
 
 func (m *Metrics) Inc(key string) {
+	m.Add(key, 1)
+}
+
+func (m *Metrics) Add(key string, value int) {
 	m.Lock()
 	defer m.Unlock()
-	m.metrics[key]++
+	m.metrics[key] += value
 }
 
 func (m *Metrics) Get(key string) int {

pkg/v2/metrics/prometheus.go

@@ -3,6 +3,7 @@ package metrics
 import (
 	"fmt"
 	"io"
+	"strings"
 )
 
 func (m *Metrics) WritePrometheus(w io.Writer) {
@@ -15,10 +16,12 @@ func (m *Metrics) WritePrometheus(w io.Writer) {
 			typ = "counter"
 		}
 
+		name := strings.SplitN(key, "{", 2)[0]
+
 		_, _ = w.Write([]byte(fmt.Sprintf(
 			"# help %s\n"+
 				"# type %s %s\n"+
 				"%s %d\n",
-			key, key, typ, key, value)))
+			name, name, typ, key, value)))
 	}
 }