diff --git a/go.mod b/go.mod index 26011f9d..f691825a 100644 --- a/go.mod +++ b/go.mod @@ -36,7 +36,7 @@ require ( github.com/Microsoft/go-winio v0.6.2 github.com/alexflint/go-arg v1.6.1 github.com/getlantern/amp v0.0.0-20260606002220-a8629924577c - github.com/getlantern/broflake v0.0.0-20261004155437-e0237743ac3d + github.com/getlantern/broflake v0.0.0-20261018030031-c6268bfc1b56 github.com/getlantern/common v1.2.1-0.20260910154003-08e030318f24 github.com/getlantern/domainfront v0.0.0-20260915172349-4984689a8667 github.com/getlantern/kindling v0.0.0-20261014210827-0f4a89dbe1ee diff --git a/go.sum b/go.sum index 3c8dbd8a..853147fb 100644 --- a/go.sum +++ b/go.sum @@ -229,8 +229,8 @@ github.com/getlantern/algeneva v0.0.0-20260731172949-952f13725b7b h1:3wQsYZSl/Uv github.com/getlantern/algeneva v0.0.0-20260731172949-952f13725b7b/go.mod h1:PrNR8tMXO26YNs8K9653XCUH7u2Kv4OdfFC3Ke1GsX0= github.com/getlantern/amp v0.0.0-20260606002220-a8629924577c h1:ZzxuhIWO295y4eE6wPYUq8Smntlgg+y39L5OOyMJY+k= github.com/getlantern/amp v0.0.0-20260606002220-a8629924577c/go.mod h1:b5teAOFT+vpBqc2CHoz74QTEM15iOv1PLdQogwXcfrM= -github.com/getlantern/broflake v0.0.0-20261004155437-e0237743ac3d h1:7P5xYoOlfp3Kd2UjNoaFTnFzc6acXvMiSK08mw0ealg= -github.com/getlantern/broflake v0.0.0-20261004155437-e0237743ac3d/go.mod h1:SVwjIH77uUEQ25MCXu+36ZnfhhFNrs1govhujY3WvYc= +github.com/getlantern/broflake v0.0.0-20261018030031-c6268bfc1b56 h1:cO2BSlKzP7C5z95b+vKH6WD8xR17X18gSOg/6RPjrBY= +github.com/getlantern/broflake v0.0.0-20261018030031-c6268bfc1b56/go.mod h1:ISVCVZlz+2T/yhWHKCmu8t8Tmw4y0+K6AxFYdcWqF70= github.com/getlantern/common v1.2.1-0.20260910154003-08e030318f24 h1:RrJsgaBsK+IspwFX7JwjFRItCWKZoKVY7UcsugheQvg= github.com/getlantern/common v1.2.1-0.20260910154003-08e030318f24/go.mod h1:9iL78jR5kZqr0E7oPxi38o6gPMhVa5dJBC5ngUgHzdg= github.com/getlantern/context v0.0.0-20220418194847-3d5e7a086201 h1:oEZYEpZo28Wdx+5FZo4aU7JFXu0WG/4wJWese5reQSA= diff --git a/unbounded/snapshot.go b/unbounded/snapshot.go index 12580628..a9964955 100644 --- a/unbounded/snapshot.go +++ b/unbounded/snapshot.go @@ -4,18 +4,106 @@ import ( "context" "sort" "strconv" + "sync" + "sync/atomic" "time" + + "github.com/getlantern/broflake/clientcore" ) var snapshotEpoch = strconv.FormatInt(time.Now().UnixNano(), 10) +// peerRetention keeps a departed consumer's byte totals in the snapshot long +// enough for the app's 1s poller to observe its final count. +const peerRetention = 60 * time.Second + // Snapshot describes the current widget state, including peers missed by event subscribers. +// Byte counts are cumulative since Epoch, like Arrivals. type Snapshot struct { - Epoch string `json:"epoch"` - Arrivals uint64 `json:"arrivals"` - Enabled bool `json:"enabled"` - Running bool `json:"running"` - Peers []string `json:"peers"` + Epoch string `json:"epoch"` + Arrivals uint64 `json:"arrivals"` + Enabled bool `json:"enabled"` + Running bool `json:"running"` + Peers []string `json:"peers"` + BytesDown uint64 `json:"bytesDown"` + BytesUp uint64 `json:"bytesUp"` + PeerBytes []PeerBytes `json:"peerBytes"` +} + +// PeerBytes is the cumulative traffic relayed for one consumer IP. Counts only +// grow while the entry exists; an entry pruned after peerRetention restarts at +// zero if the IP returns, which consumers treat as a reset. +type PeerBytes struct { + IP string `json:"ip"` + Down uint64 `json:"down"` + Up uint64 `json:"up"` +} + +// slotBytes holds bytes relayed through one broflake worker slot since they +// were last folded into the manager's totals. Chunk callbacks only touch these +// atomics, so the per-chunk data path never contends on m.mu. +type slotBytes struct { + down, up atomic.Uint64 +} + +// slotTable is created per worker run so chunks a stopped worker delivers late +// land in an orphaned table instead of the next run's attribution. +type slotTable struct { + slots sync.Map + // mu is read-held by each chunk callback around its add, so close can + // wait out callbacks already past the closed check before the final fold. + mu sync.RWMutex + closed bool +} + +func (t *slotTable) add(idx int, size uint64, down bool) { + t.mu.RLock() + defer t.mu.RUnlock() + if t.closed { + return + } + s := t.slot(idx) + if down { + s.down.Add(size) + } else { + s.up.Add(size) + } +} + +// close stops the table accepting bytes and returns once every in-flight add +// has landed, so a fold after close sees all accepted bytes. +func (t *slotTable) close() { + t.mu.Lock() + t.closed = true + t.mu.Unlock() +} + +func (t *slotTable) slot(idx int) *slotBytes { + if s, ok := t.slots.Load(idx); ok { + return s.(*slotBytes) + } + s, _ := t.slots.LoadOrStore(idx, &slotBytes{}) + return s.(*slotBytes) +} + +type peerTally struct { + down, up uint64 + // active counts slots currently held by this IP; endedAt is set when it drops to zero. + active int + endedAt time.Time +} + +// chunkCounter returns a broflake ChunkFunc that adds relayed bytes to the +// slot's downstream or upstream counter. It keeps counting after the worker's +// context is cancelled: broflake tears down asynchronously after ui.Stop +// returns and can still deliver bytes, which the retired run absorbs. +func chunkCounter(slots *slotTable, down bool) clientcore.ChunkFunc { + return func(size int, workerIdx int) { + if size <= 0 { + return + } + slots.add(workerIdx, uint64(size), down) + } } // CurrentSnapshot returns an independent copy of the live widget state. @@ -26,7 +114,17 @@ func CurrentSnapshot() Snapshot { func (m *unboundedManager) snapshot() Snapshot { m.mu.Lock() defer m.mu.Unlock() - s := Snapshot{Epoch: snapshotEpoch, Arrivals: m.arrivals, Enabled: Enabled(), Running: m.running, Peers: []string{}} + m.prunePeerBytes() + s := Snapshot{ + Epoch: snapshotEpoch, + Arrivals: m.arrivals, + Enabled: Enabled(), + Running: m.running, + Peers: []string{}, + BytesDown: m.bytesDown, + BytesUp: m.bytesUp, + PeerBytes: []PeerBytes{}, + } if m.running { for _, source := range m.peers { if source != "" { @@ -35,6 +133,10 @@ func (m *unboundedManager) snapshot() Snapshot { } } sort.Strings(s.Peers) + for ip, t := range m.peerBytes { + s.PeerBytes = append(s.PeerBytes, PeerBytes{IP: ip, Down: t.down, Up: t.up}) + } + sort.Slice(s.PeerBytes, func(i, j int) bool { return s.PeerBytes[i].IP < s.PeerBytes[j].IP }) return s } @@ -44,6 +146,9 @@ func (m *unboundedManager) recordConnection(ctx context.Context, state, slot int if ctx.Err() != nil { return } + // Prune before acquiring, so a returning IP's fresh tally can't pick up + // bytes still pending from its pruned lifetime. + m.prunePeerBytes() if state > 0 { if m.peers == nil { m.peers = make(map[int]string) @@ -60,8 +165,173 @@ func (m *unboundedManager) recordConnection(ctx context.Context, state, slot int m.arrivals++ } } + // Bytes counted before this accept belong to the slot's previous occupant. + m.foldSlot(slot) + if prev, ok := m.peers[slot]; ok { + m.releasePeer(prev) + } + delete(m.lastOwner, slot) m.peers[slot] = source + m.acquirePeer(source) } else if state < 0 { + m.foldSlot(slot) + if prev, ok := m.peers[slot]; ok { + m.releasePeer(prev) + // broflake batches upstream bytes for up to 250ms and reports the + // close without flushing them, so the closed peer keeps the slot's + // bytes until the next accept. + if m.lastOwner == nil { + m.lastOwner = make(map[int]*peerTally) + } + m.lastOwner[slot] = m.peerBytes[prev] + } delete(m.peers, slot) } } + +// retiredRun keeps a stopped run's slot table and its final slot owners, so +// bytes broflake delivers during its asynchronous teardown are still counted +// and credited to the consumers that relayed them. +type retiredRun struct { + slots *slotTable + owners map[int]*peerTally + retiredAt time.Time +} + +// releaseAllPeers folds outstanding slot bytes, ends every live peer so their +// final totals stay visible for peerRetention, and retires the run's slot +// table for the same period. Caller must hold m.mu. +func (m *unboundedManager) releaseAllPeers() { + m.prunePeerBytes() + for _, source := range m.peers { + m.releasePeer(source) + } + if m.slots != nil { + owners := make(map[int]*peerTally, len(m.lastOwner)+len(m.peers)) + for idx, t := range m.lastOwner { + owners[idx] = t + } + for idx, source := range m.peers { + owners[idx] = m.peerBytes[source] + } + m.retired = append(m.retired, retiredRun{slots: m.slots, owners: owners, retiredAt: m.clock()}) + } + m.slots = nil + m.lastOwner = nil +} + +// foldRetiredRuns folds bytes stopped runs received since the last fold, and +// drops runs retired longer than peerRetention after a final fold. Caller +// must hold m.mu. +func (m *unboundedManager) foldRetiredRuns() { + now := m.clock() + kept := m.retired[:0] + for _, r := range m.retired { + expired := now.Sub(r.retiredAt) > peerRetention + if expired { + r.slots.close() + } + r.slots.slots.Range(func(k, v any) bool { + m.addSlotBytes(r.owners[k.(int)], v.(*slotBytes)) + return true + }) + if !expired { + kept = append(kept, r) + } + } + clear(m.retired[len(kept):]) + m.retired = kept +} + +// foldSlot moves a slot's pending bytes into the totals and, when the slot's +// IP is known, that IP's tally. Caller must hold m.mu. +func (m *unboundedManager) foldSlot(idx int) { + if m.slots == nil { + return + } + if s, ok := m.slots.slots.Load(idx); ok { + m.addSlotBytes(m.slotOwner(idx), s.(*slotBytes)) + } +} + +func (m *unboundedManager) foldAllSlots() { + if m.slots == nil { + return + } + m.slots.slots.Range(func(k, v any) bool { + m.addSlotBytes(m.slotOwner(k.(int)), v.(*slotBytes)) + return true + }) +} + +// slotOwner is the tally credited with a live slot's bytes: its current +// consumer's, or the last one's while the slot is vacant. Owners are held as +// tallies rather than IPs, so bytes from a pruned lifetime never reach a +// returning IP's fresh tally. +func (m *unboundedManager) slotOwner(idx int) *peerTally { + if owner, live := m.peers[idx]; live { + return m.peerBytes[owner] + } + return m.lastOwner[idx] +} + +func (m *unboundedManager) addSlotBytes(t *peerTally, s *slotBytes) { + down, up := s.down.Swap(0), s.up.Swap(0) + if down == 0 && up == 0 { + return + } + m.bytesDown += down + m.bytesUp += up + if t != nil { + t.down += down + t.up += up + } +} + +func (m *unboundedManager) acquirePeer(source string) { + if source == "" { + return + } + if m.peerBytes == nil { + m.peerBytes = make(map[string]*peerTally) + } + t := m.peerBytes[source] + if t == nil { + t = &peerTally{} + m.peerBytes[source] = t + } + t.active++ +} + +func (m *unboundedManager) releasePeer(source string) { + t := m.peerBytes[source] + if t == nil || t.active == 0 { + return + } + t.active-- + if t.active == 0 { + t.endedAt = m.clock() + } +} + +// prunePeerBytes drops tallies of IPs ended more than peerRetention ago, and +// expired retired runs. It first folds every pending byte, so bytes are +// always credited to the tally lifetime they were relayed in. Caller must +// hold m.mu. +func (m *unboundedManager) prunePeerBytes() { + m.foldAllSlots() + m.foldRetiredRuns() + now := m.clock() + for ip, t := range m.peerBytes { + if t.active == 0 && now.Sub(t.endedAt) > peerRetention { + delete(m.peerBytes, ip) + } + } +} + +func (m *unboundedManager) clock() time.Time { + if m.now != nil { + return m.now() + } + return time.Now() +} diff --git a/unbounded/snapshot_test.go b/unbounded/snapshot_test.go index 4cd4a119..d08bc663 100644 --- a/unbounded/snapshot_test.go +++ b/unbounded/snapshot_test.go @@ -2,6 +2,8 @@ package unbounded import ( "context" + "encoding/json" + "sync" "testing" "time" @@ -51,3 +53,368 @@ func TestSnapshotReportsActualLifecycle(t *testing.T) { require.False(t, CurrentSnapshot().Running) require.Empty(t, CurrentSnapshot().Peers) } + +func newBytesManager(now *time.Time) (*unboundedManager, *slotTable) { + slots := &slotTable{} + return &unboundedManager{running: true, slots: slots, now: func() time.Time { return *now }}, slots +} + +func peerBytesByIP(s Snapshot) map[string]PeerBytes { + out := make(map[string]PeerBytes) + for _, p := range s.PeerBytes { + out[p.IP] = p + } + return out +} + +func TestSnapshotAttributesBytesToPeerIPs(t *testing.T) { + now := time.Unix(1000, 0) + m, slots := newBytesManager(&now) + ctx := context.Background() + down, up := chunkCounter(slots, true), chunkCounter(slots, false) + + m.recordConnection(ctx, 1, 0, "192.0.2.1") + m.recordConnection(ctx, 1, 1, "192.0.2.2") + m.recordConnection(ctx, 1, 2, "192.0.2.1") + down(100, 0) + down(50, 2) + up(7, 0) + down(30, 1) + up(3, 1) + // No known IP for slot 9, so these count only toward the totals. + down(1000, 9) + up(1, 9) + down(-5, 0) + + s := m.snapshot() + require.Equal(t, uint64(1180), s.BytesDown) + require.Equal(t, uint64(11), s.BytesUp) + require.Equal(t, []PeerBytes{ + {IP: "192.0.2.1", Down: 150, Up: 7}, + {IP: "192.0.2.2", Down: 30, Up: 3}, + }, s.PeerBytes) + + down(10, 1) + s = m.snapshot() + require.Equal(t, uint64(1190), s.BytesDown) + require.Equal(t, uint64(40), peerBytesByIP(s)["192.0.2.2"].Down) +} + +func TestSnapshotSlotReuseByDifferentIP(t *testing.T) { + now := time.Unix(1000, 0) + m, slots := newBytesManager(&now) + ctx := context.Background() + down := chunkCounter(slots, true) + + m.recordConnection(ctx, 1, 0, "192.0.2.1") + down(100, 0) + m.recordConnection(ctx, -1, 0, "192.0.2.1") + // Chunks between close and the next accept are late batches from the closed peer. + down(5, 0) + m.recordConnection(ctx, 1, 0, "192.0.2.2") + down(40, 0) + // An accept without a preceding close still hands the slot over cleanly. + m.recordConnection(ctx, 1, 0, "192.0.2.3") + down(2, 0) + + s := m.snapshot() + require.Equal(t, uint64(147), s.BytesDown) + got := peerBytesByIP(s) + require.Equal(t, uint64(105), got["192.0.2.1"].Down) + require.Equal(t, uint64(40), got["192.0.2.2"].Down) + require.Equal(t, uint64(2), got["192.0.2.3"].Down) + require.Equal(t, 0, m.peerBytes["192.0.2.2"].active) + require.Equal(t, 1, m.peerBytes["192.0.2.3"].active) +} + +func TestSnapshotRetainsEndedPeersFor60s(t *testing.T) { + now := time.Unix(1000, 0) + m, slots := newBytesManager(&now) + ctx := context.Background() + down := chunkCounter(slots, true) + + m.recordConnection(ctx, 1, 0, "192.0.2.1") + m.recordConnection(ctx, 1, 1, "192.0.2.1") + m.recordConnection(ctx, 1, 2, "192.0.2.2") + down(10, 0) + down(20, 2) + m.recordConnection(ctx, -1, 0, "") + m.recordConnection(ctx, -1, 2, "") + + now = now.Add(2 * peerRetention) + // 192.0.2.1 still holds slot 1, so it's kept regardless of age. + require.Len(t, m.snapshot().PeerBytes, 1) + m.recordConnection(ctx, -1, 1, "") + + now = now.Add(peerRetention) + s := m.snapshot() + require.Equal(t, []PeerBytes{{IP: "192.0.2.1", Down: 10}}, s.PeerBytes) + now = now.Add(time.Second) + s = m.snapshot() + require.Empty(t, s.PeerBytes) + require.Equal(t, uint64(30), s.BytesDown) + + // A returning IP starts a fresh entry. + m.recordConnection(ctx, 1, 0, "192.0.2.1") + down(4, 0) + require.Equal(t, []PeerBytes{{IP: "192.0.2.1", Down: 4}}, m.snapshot().PeerBytes) +} + +func TestSnapshotStopKeepsFinalPeerBytes(t *testing.T) { + resetManager(t, func(*clientcore.BroflakeOptions, *clientcore.WebRTCOptions, *clientcore.EgressOptions) (widget, error) { + return &fakeWidget{}, nil + }) + now := time.Unix(1000, 0) + manager.now = func() time.Time { return now } + require.NoError(t, settings.Set(settings.UnboundedKey, true)) + manager.mu.Lock() + manager.lastFeatureOn = true + manager.lastCfg = testCfg() + manager.mu.Unlock() + require.NoError(t, Apply()) + require.Eventually(t, func() bool { return CurrentSnapshot().Running }, time.Second, time.Millisecond) + + manager.mu.Lock() + slots := manager.slots + manager.mu.Unlock() + ctx, stopCtx := context.WithCancel(context.Background()) + manager.recordConnection(ctx, 1, 0, "192.0.2.1") + down := chunkCounter(slots, true) + down(64, 0) + + require.NoError(t, Stop(context.Background())) + stopCtx() + s := CurrentSnapshot() + require.False(t, s.Running) + require.Empty(t, s.Peers) + require.Equal(t, uint64(64), s.BytesDown) + require.Equal(t, []PeerBytes{{IP: "192.0.2.1", Down: 64}}, s.PeerBytes) + + // broflake tears down asynchronously, so chunks can arrive after Stop. + down(36, 0) + require.Equal(t, []PeerBytes{{IP: "192.0.2.1", Down: 100}}, CurrentSnapshot().PeerBytes) + + now = now.Add(peerRetention + time.Second) + s = CurrentSnapshot() + require.Empty(t, s.PeerBytes) + require.Equal(t, uint64(100), s.BytesDown) + manager.mu.Lock() + require.Empty(t, manager.retired) + manager.mu.Unlock() + // Once the retired run expires its table is closed. + down(1000, 0) + require.Equal(t, uint64(100), CurrentSnapshot().BytesDown) +} + +func TestSnapshotJSONShape(t *testing.T) { + now := time.Unix(1000, 0) + m, slots := newBytesManager(&now) + ctx := context.Background() + m.recordConnection(ctx, 1, 0, "192.0.2.1") + chunkCounter(slots, true)(12, 0) + chunkCounter(slots, false)(3, 0) + + raw, err := json.Marshal(m.snapshot()) + require.NoError(t, err) + var got map[string]any + require.NoError(t, json.Unmarshal(raw, &got)) + require.ElementsMatch(t, []string{"epoch", "arrivals", "enabled", "running", "peers", "bytesDown", "bytesUp", "peerBytes"}, keys(got)) + require.Equal(t, float64(12), got["bytesDown"]) + require.Equal(t, float64(3), got["bytesUp"]) + require.Equal(t, []any{map[string]any{"ip": "192.0.2.1", "down": float64(12), "up": float64(3)}}, got["peerBytes"]) + + raw, err = json.Marshal((&unboundedManager{}).snapshot()) + require.NoError(t, err) + require.Contains(t, string(raw), `"peerBytes":[]`) + + // Payloads from older backends without byte fields still decode. + var old Snapshot + require.NoError(t, json.Unmarshal([]byte(`{"epoch":"1","arrivals":2,"enabled":true,"running":true,"peers":[]}`), &old)) + require.Zero(t, old.BytesDown) + require.Empty(t, old.PeerBytes) +} + +func keys(m map[string]any) []string { + out := make([]string, 0, len(m)) + for k := range m { + out = append(out, k) + } + return out +} + +func TestSnapshotConcurrentChunksAreCountedOnce(t *testing.T) { + now := time.Unix(1000, 0) + m, slots := newBytesManager(&now) + ctx := context.Background() + down := chunkCounter(slots, true) + m.recordConnection(ctx, 1, 0, "192.0.2.1") + m.recordConnection(ctx, 1, 1, "192.0.2.2") + + var wg sync.WaitGroup + for w := 0; w < 2; w++ { + wg.Add(1) + go func(w int) { + defer wg.Done() + for i := 0; i < 1000; i++ { + down(1, w) + } + }(w) + } + for i := 0; i < 50; i++ { + m.snapshot() + } + wg.Wait() + s := m.snapshot() + require.Equal(t, uint64(2000), s.BytesDown) + require.Equal(t, []PeerBytes{{IP: "192.0.2.1", Down: 1000}, {IP: "192.0.2.2", Down: 1000}}, s.PeerBytes) +} + +func TestSnapshotCreditsLateUpstreamBatchToClosedPeer(t *testing.T) { + now := time.Unix(1000, 0) + m, slots := newBytesManager(&now) + ctx := context.Background() + up := chunkCounter(slots, false) + + m.recordConnection(ctx, 1, 0, "192.0.2.1") + up(10, 0) + m.snapshot() + m.recordConnection(ctx, -1, 0, "") + m.snapshot() + // broflake flushes the closed peer's pending upstream batch after the close. + up(7, 0) + now = now.Add(time.Second) + require.Equal(t, []PeerBytes{{IP: "192.0.2.1", Up: 17}}, m.snapshot().PeerBytes) + + m.recordConnection(ctx, 1, 0, "192.0.2.2") + up(3, 0) + s := m.snapshot() + require.Equal(t, uint64(20), s.BytesUp) + got := peerBytesByIP(s) + require.Equal(t, uint64(17), got["192.0.2.1"].Up) + require.Equal(t, uint64(3), got["192.0.2.2"].Up) +} + +func TestSlotTableCloseWaitsForInFlightAdds(t *testing.T) { + slots := &slotTable{} + slots.mu.RLock() + closed := make(chan struct{}) + go func() { + slots.close() + close(closed) + }() + select { + case <-closed: + t.Fatal("close returned while an add was in flight") + case <-time.After(20 * time.Millisecond): + } + slots.slot(0).down.Add(5) + slots.mu.RUnlock() + <-closed + slots.add(0, 100, true) + require.Equal(t, uint64(5), slots.slot(0).down.Load()) +} + +type flushingWidget struct{ onStop func() } + +func (w *flushingWidget) Stop() { w.onStop() } + +func TestSnapshotCountsUpstreamFlushedWhileStopping(t *testing.T) { + var upstream clientcore.ChunkFunc + resetManager(t, func(bf *clientcore.BroflakeOptions, _ *clientcore.WebRTCOptions, _ *clientcore.EgressOptions) (widget, error) { + upstream = bf.OnUpstreamChunkFunc + return &flushingWidget{onStop: func() { upstream(9, 0) }}, nil + }) + now := time.Unix(1000, 0) + manager.now = func() time.Time { return now } + require.NoError(t, settings.Set(settings.UnboundedKey, true)) + manager.mu.Lock() + manager.lastFeatureOn = true + manager.lastCfg = testCfg() + manager.mu.Unlock() + require.NoError(t, Apply()) + require.Eventually(t, func() bool { return CurrentSnapshot().Running }, time.Second, time.Millisecond) + + manager.recordConnection(context.Background(), 1, 0, "192.0.2.1") + upstream(4, 0) + require.NoError(t, Stop(context.Background())) + s := CurrentSnapshot() + require.Equal(t, uint64(13), s.BytesUp) + require.Equal(t, []PeerBytes{{IP: "192.0.2.1", Up: 13}}, s.PeerBytes) +} + +func TestSnapshotCreditsRetiredRunBytesToItsOwners(t *testing.T) { + now := time.Unix(1000, 0) + m, oldSlots := newBytesManager(&now) + ctx := context.Background() + oldUp := chunkCounter(oldSlots, false) + m.recordConnection(ctx, 1, 0, "192.0.2.1") + m.recordConnection(ctx, 1, 1, "192.0.2.2") + m.recordConnection(ctx, -1, 1, "") + m.mu.Lock() + m.releaseAllPeers() + newSlots := &slotTable{} + m.slots = newSlots + m.peers = make(map[int]string) + m.mu.Unlock() + + newUp := chunkCounter(newSlots, false) + m.recordConnection(ctx, 1, 0, "192.0.2.3") + oldUp(5, 0) + oldUp(2, 1) + newUp(3, 0) + s := m.snapshot() + require.Equal(t, uint64(10), s.BytesUp) + require.Equal(t, []PeerBytes{ + {IP: "192.0.2.1", Up: 5}, + {IP: "192.0.2.2", Up: 2}, + {IP: "192.0.2.3", Up: 3}, + }, s.PeerBytes) +} + +func TestSnapshotPendingBytesStayWithPrunedLifetime(t *testing.T) { + now := time.Unix(1000, 0) + m, slots := newBytesManager(&now) + ctx := context.Background() + up := chunkCounter(slots, false) + + m.recordConnection(ctx, 1, 0, "192.0.2.1") + up(5, 0) + m.recordConnection(ctx, -1, 0, "") + now = now.Add(peerRetention + time.Second) + // A late batch is still pending when the same IP returns, with no + // snapshot taken in between. + up(7, 0) + m.recordConnection(ctx, 1, 1, "192.0.2.1") + s := m.snapshot() + require.Equal(t, uint64(12), s.BytesUp) + require.Equal(t, []PeerBytes{{IP: "192.0.2.1"}}, s.PeerBytes) +} + +func TestRetiredRunsExpireWithoutPolling(t *testing.T) { + now := time.Unix(1000, 0) + m, _ := newBytesManager(&now) + m.mu.Lock() + defer m.mu.Unlock() + m.releaseAllPeers() + m.slots = &slotTable{} + now = now.Add(peerRetention + time.Second) + m.releaseAllPeers() + require.Len(t, m.retired, 1) +} + +func TestSnapshotLateBytesAfterPruneSkipReturningIP(t *testing.T) { + now := time.Unix(1000, 0) + m, slots := newBytesManager(&now) + ctx := context.Background() + up := chunkCounter(slots, false) + + m.recordConnection(ctx, 1, 0, "192.0.2.1") + m.recordConnection(ctx, -1, 0, "") + now = now.Add(peerRetention + time.Second) + m.recordConnection(ctx, 1, 1, "192.0.2.1") + // Arrives after the prune that started the IP's new lifetime. + up(7, 0) + s := m.snapshot() + require.Equal(t, uint64(7), s.BytesUp) + require.Equal(t, []PeerBytes{{IP: "192.0.2.1"}}, s.PeerBytes) +} diff --git a/unbounded/unbounded.go b/unbounded/unbounded.go index 0bd7f7c0..e466ecc9 100644 --- a/unbounded/unbounded.go +++ b/unbounded/unbounded.go @@ -183,6 +183,21 @@ type unboundedManager struct { running bool peers map[int]string arrivals uint64 + + // slots receives per-chunk byte counts from the live worker; it is + // folded into bytesDown/bytesUp and peerBytes (keyed by consumer IP) + // whenever a slot's connection changes or a snapshot is taken. + slots *slotTable + bytesDown uint64 + bytesUp uint64 + peerBytes map[string]*peerTally + // lastOwner keeps a closed slot's tally so bytes broflake reports after the + // close are still credited to it. + lastOwner map[int]*peerTally + // retired holds stopped runs' tables still absorbing late bytes. + retired []retiredRun + // now overrides time.Now for peer retention tests. + now func() time.Time } // shouldStart reports whether all three start conditions hold. Caller @@ -439,6 +454,7 @@ func (m *unboundedManager) stopCtx(ctx context.Context, disarm bool) error { cancel() m.mu.Lock() m.running = false + m.releaseAllPeers() m.peers = nil m.mu.Unlock() select { @@ -498,6 +514,8 @@ func (m *unboundedManager) start() { // Retain the immutable config snapshot to detect changes while this worker runs. m.runningCfg = ucfg m.peers = make(map[int]string) + slots := &slotTable{} + m.slots = slots m.mu.Unlock() go func() { @@ -550,6 +568,13 @@ func (m *unboundedManager) start() { }) } + // Downstream chunks fire per chunk on broflake's data path, so the + // counters stay lock-free and attribution to an IP happens later + // under m.mu. Chunks arriving after stop land in the retired run, + // credited to the slot owners at stop. + bfOpt.OnDownstreamChunkFunc = chunkCounter(slots, true) + bfOpt.OnUpstreamChunkFunc = chunkCounter(slots, false) + rtcOpt := clientcore.NewDefaultWebRTCOptions() rtcOpt.STUNBatch = donorSTUNBatch(nil) if ucfg != nil {