bug(mq): an owner cut off from NATS keeps inserting its held rows for up to a minute after the next owner starts on the same tables #705

package aud

import (
	"bufio"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/Wave-RF/WaveHouse/internal/chconn"
	"github.com/Wave-RF/WaveHouse/internal/ingest"
	"github.com/Wave-RF/WaveHouse/internal/mq"
	"github.com/Wave-RF/WaveHouse/internal/tenant"
)

// blackhole is a TCP proxy that can stop moving bytes in both directions
// without closing anything: a network partition the client only notices by
// its own timeouts.
type blackhole struct {
	ln     net.Listener
	paused atomic.Bool
}

func newBlackhole(t *testing.T, listen, upstream string) *blackhole {
	t.Helper()
	ln, err := net.Listen("tcp", listen)
	require.NoError(t, err)
	b := &blackhole{ln: ln}
	t.Cleanup(func() { _ = ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			u, err := net.Dial("tcp", upstream)
			if err != nil {
				_ = c.Close()
				continue
			}
			go b.pipe(c, u)
			go b.pipe(u, c)
		}
	}()
	return b
}

func (b *blackhole) pipe(dst, src net.Conn) {
	buf := make([]byte, 32<<10)
	for {
		for b.paused.Load() {
			time.Sleep(20 * time.Millisecond)
		}
		n, err := src.Read(buf)
		for b.paused.Load() { // hold what was read until the heal
			time.Sleep(20 * time.Millisecond)
		}
		if n > 0 {
			if _, werr := dst.Write(buf[:n]); werr != nil {
				return
			}
		}
		if err != nil {
			if err != io.EOF {
				_ = dst.Close()
			}
			return
		}
	}
}

type write struct {
	proc string
	at   time.Time
	seqs []int
}

// sharedCH is one ClickHouse both processes write to, taking latency per
// insert; it records who wrote which rows of which table, when.
type sharedCH struct {
	mu      sync.Mutex
	writes  map[string][]write
	latency time.Duration
}

func (c *sharedCH) server(t *testing.T, proc string) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		table := r.URL.Query().Get("param_target_table")
		var seqs []int
		sc := bufio.NewScanner(r.Body)
		for sc.Scan() {
			var row []int
			_ = json.Unmarshal(sc.Bytes(), &row)
			seqs = append(seqs, row[0])
		}
		time.Sleep(c.latency)
		c.mu.Lock()
		if c.writes == nil {
			c.writes = map[string][]write{}
		}
		c.writes[table] = append(c.writes[table], write{proc: proc, at: time.Now(), seqs: seqs})
		c.mu.Unlock()
		w.WriteHeader(http.StatusOK)
	}))
	t.Cleanup(srv.Close)
	return srv
}

// E2: the owner of some shards (A) is cut off from NATS, but not from
// ClickHouse, while it holds rows. It cannot renew its lease or its pins, so
// B takes its shards over after the pin TTL and resets them. Does A keep
// writing rows after B has started writing the same tables (no self-stop on
// the write path)? What does that do to per-table first-write order and
// duplicates?
func TestAudit_E2_PartitionedOwnerKeepsWriting(t *testing.T) {
	const tables = 8
	srv := startServer(t, envInt("AUD_PORT", 57111))
	bh := newBlackhole(t, "127.0.0.1:"+strconv.Itoa(envInt("AUD_PROXY", 57112)), srv.s.ClientURL()[len("nats://"):])
	pub := broker(t, srv.s.ClientURL())
	ch := &sharedCH{latency: time.Duration(envInt("AUD_CH_MS", 1500)) * time.Millisecond}
	chA, chB := ch.server(t, "A"), ch.server(t, "B")
	targetOf := func(u string) func(tenant.ID) chconn.Target {
		return func(tenant.ID) chconn.Target { return chconn.Target{URL: u, Database: "db"} }
	}
	cfg := ingest.ClaimConfig{Every: 100 * time.Millisecond}
	lease := 3 * time.Second

	// Publish continuously: tables t0..t7, 100 rows/s each.
	stopPub := make(chan struct{})
	var pubWG sync.WaitGroup
	next := make([]int, tables)
	pubWG.Go(func() {
		tick := time.NewTicker(10 * time.Millisecond)
		defer tick.Stop()
		for {
			select {
			case <-stopPub:
				return
			case <-tick.C:
			}
			for i := range tables {
				tb := "t" + strconv.Itoa(i)
				if err := pub.Publish(t.Context(), mq.Topic{Tenant: "acme", Table: tb}, envelope(t, tb, next[i])); err == nil {
					next[i]++
				}
			}
		}
	})

	a := startProc(t, "nats://"+bh.ln.Addr().String(), "A", lease, cfg, targetOf(chA.URL))
	time.Sleep(4 * time.Second)
	b := startProc(t, srv.s.ClientURL(), "B", lease, cfg, targetOf(chB.URL))
	time.Sleep(12 * time.Second) // rebalance settles; A keeps some tables

	ch.mu.Lock()
	lastWriter := map[string]string{}
	for tb, ws := range ch.writes {
		lastWriter[tb] = ws[len(ws)-1].proc
	}
	ch.mu.Unlock()
	t.Logf("last writer per table before the partition: %v", lastWriter)

	partAt := time.Now()
	bh.paused.Store(true)
	t.Logf("partitioned A from NATS")
	time.Sleep(time.Duration(envInt("AUD_PART_S", 45)) * time.Second)
	bh.paused.Store(false)
	healAt := time.Now()
	t.Logf("healed")
	time.Sleep(20 * time.Second)
	close(stopPub)
	pubWG.Wait()
	// Let everything drain.
	time.Sleep(15 * time.Second)

	ch.mu.Lock()
	defer ch.mu.Unlock()
	for i := range tables {
		tb := "t" + strconv.Itoa(i)
		ws := ch.writes[tb]
		var all []int
		var firstB time.Time
		for _, w := range ws {
			all = append(all, w.seqs...)
			if w.proc == "B" && w.at.After(partAt) && firstB.IsZero() {
				firstB = w.at
			}
		}
		var staleRows, staleWrites int
		var lastA time.Time
		for _, w := range ws {
			if w.proc == "A" && w.at.After(partAt) {
				lastA = w.at
				if !firstB.IsZero() && w.at.After(firstB) {
					staleWrites++
					staleRows += len(w.seqs)
				}
			}
		}
		if lastWriter[tb] == "A" {
			for _, w := range ws {
				if w.at.After(partAt.Add(-3*time.Second)) && w.at.Before(healAt.Add(15*time.Second)) && len(w.seqs) > 0 {
					t.Logf("  %s %s +%s rows %d..%d (%d)", tb, w.proc, w.at.Sub(partAt).Round(100*time.Millisecond), w.seqs[0], w.seqs[len(w.seqs)-1], len(w.seqs))
				}
			}
			t.Logf("  heal at +%s", healAt.Sub(partAt).Round(100*time.Millisecond))
		}
		inv, dups, distinct := orderStats(all)
		t.Logf("RESULT %s owner-before=%s published=%d distinct=%d dups=%d first-write-inversions=%d | A's last write %s after partition; B's first %s after; A writes after B's first: %d inserts, %d rows",
			tb, lastWriter[tb], next[i], distinct, dups, inv,
			since(partAt, lastA), since(partAt, firstB), staleWrites, staleRows)
	}
	for _, p := range []*proc{a, b} {
		select {
		case err := <-p.failed:
			t.Logf("%s failed: %v", p.id, err)
		default:
		}
	}
}

func since(from, at time.Time) string {
	if at.IsZero() {
		return "never"
	}
	return at.Sub(from).Round(100 * time.Millisecond).String()
}