bastion 0.61.2

client-go/lanehold.go

@@ -0,0 +1,157 @@
+package main
+
+// When a Host holds this machine (its owner quarantined it there), the Host
+// refuses every worker socket of every lane with close code 4403, and says
+// "quarantined": true on each beat it answers. The lane workers park a pool of
+// sockets per lane and used to dial again at once (mgmt, harness, session) or
+// a second after a close (egress, serve), so one held machine whose tunnel
+// stayed up dialled the Host 30 to 80 times a second, each one a TLS
+// handshake. A pairing's workers now share one hold:
+//
+//   - a socket the Host closes with 4403 holds every worker of that pairing for
+//     laneHoldFirst, doubling with each hold up to laneHoldMax;
+//   - a beat that says quarantined stops them dialling until a beat no longer
+//     does, which is how the owner's release reaches this machine.
+//
+// The hold is the pairing's, so another Host this machine is paired with is
+// never slowed by it.
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"math/rand/v2"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// laneHeldCloseCode is the close code a Host gives a worker socket of a
+// machine it holds (the Host's quarantine gate).
+const laneHeldCloseCode = 4403
+
+// The first hold, and the longest. Variables so a test can shorten them.
+var (
+	laneHoldFirst = 5 * time.Second
+	laneHoldMax   = 5 * time.Minute
+)
+
+type laneHold struct {
+	mu          sync.Mutex
+	wait        time.Duration // the last hold's length; the next one doubles it
+	until       time.Time     // no worker of the pairing dials before this
+	last        time.Time     // when the last hold began
+	quarantined bool          // what the last beat said
+	wake        chan struct{} // closed when the beat lifts a quarantine
+}
+
+var laneHolds sync.Map // host name -> *laneHold
+
+func laneHoldFor(host string) *laneHold {
+	if h, ok := laneHolds.Load(host); ok {
+		return h.(*laneHold)
+	}
+	h, _ := laneHolds.LoadOrStore(host, &laneHold{wake: make(chan struct{})})
+	return h.(*laneHold)
+}
+
+// refused is a worker socket the Host closed with 4403. The workers of a pool
+// are refused together, so a refusal while a hold is on is that hold's and
+// does not lengthen it. A refusal long after the last hold starts again from
+// laneHoldFirst.
+func (h *laneHold) refused(host string) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	now := time.Now()
+	if now.Before(h.until) {
+		return
+	}
+	switch {
+	case h.wait == 0 || now.Sub(h.last) > 2*laneHoldMax:
+		h.wait = laneHoldFirst
+	default:
+		h.wait = min(2*h.wait, laneHoldMax)
+	}
+	h.last = now
+	h.until = now.Add(h.wait)
+	fmt.Printf("[%s] the Host refused a lane: it holds this machine; its lanes dial again in %s\n", host, h.wait)
+}
+
+// fromBeat is what a beat's reply said. A release ends the hold at once and
+// starts the next one from laneHoldFirst.
+func (h *laneHold) fromBeat(host string, quarantined bool) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	switch {
+	case quarantined && !h.quarantined:
+		h.quarantined = true
+		fmt.Printf("[%s] the Host holds this machine (quarantined there): its lanes wait until its owner releases it\n", host)
+	case !quarantined && h.quarantined:
+		h.quarantined = false
+		h.wait, h.until = 0, time.Time{}
+		close(h.wake)
+		h.wake = make(chan struct{})
+		fmt.Printf("[%s] the Host released this machine: its lanes are back\n", host)
+	}
+}
+
+// waitClear blocks a worker until it may dial: the hold is over and the last
+// beat did not say quarantined. A worker waits a little past the hold's end,
+// a different little for each, so a pool does not dial in one burst. False
+// when ctx ends first.
+func (h *laneHold) waitClear(ctx context.Context) bool {
+	for {
+		h.mu.Lock()
+		quarantined, until, wake := h.quarantined, h.until, h.wake
+		h.mu.Unlock()
+		if ctx.Err() != nil {
+			return false
+		}
+		var t *time.Timer
+		var timer <-chan time.Time
+		switch {
+		case quarantined:
+			// Only a beat lifts it.
+		case time.Now().Before(until):
+			d := time.Until(until)
+			t = time.NewTimer(d + rand.N(d/10+time.Millisecond))
+			timer = t.C
+		default:
+			return true
+		}
+		select {
+		case <-ctx.Done():
+		case <-wake:
+		case <-timer:
+		}
+		if t != nil {
+			t.Stop()
+		}
+	}
+}
+
+// watchHeldClose reports, once the socket has been read to its end, whether
+// the Host closed it with laneHeldCloseCode.
+func watchHeldClose(ws *websocket.Conn) (held func() bool) {
+	var got atomic.Bool
+	def := ws.CloseHandler()
+	ws.SetCloseHandler(func(code int, text string) error {
+		if code == laneHeldCloseCode {
+			got.Store(true)
+		}
+		return def(code, text)
+	})
+	return got.Load
+}
+
+// beatSaysQuarantined reads the "quarantined" a Host puts on the beat's reply
+// for a machine it holds. A reply without it, or one that cannot be read, is
+// not a hold.
+func beatSaysQuarantined(body []byte) bool {
+	var r struct {
+		Quarantined bool `json:"quarantined"`
+	}
+	return json.Unmarshal(body, &r) == nil && r.Quarantined
+}