snapshot of https://github.com/nats-io/nats-server.git (main, 2026-09-15; blobs > 131072B omitted)

server/thw/helper_test.go

@@ -0,0 +1,57 @@
+// Copyright 2024 The NATS Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package thw
+
+import (
+	"strings"
+	"testing"
+)
+
+func require_NoError(t *testing.T, err error) {
+	t.Helper()
+	if err != nil {
+		t.Fatalf("require no error, but got: %v", err)
+	}
+}
+
+func require_Error(t *testing.T, err error, expected ...error) {
+	t.Helper()
+	if err == nil {
+		t.Fatalf("require error, but got none")
+	}
+	if len(expected) == 0 {
+		return
+	}
+	eStr := err.Error()
+	for _, e := range expected {
+		if err == e || strings.Contains(eStr, e.Error()) || strings.Contains(e.Error(), eStr) {
+			return
+		}
+	}
+	t.Fatalf("Expected one of %v, got '%v'", expected, err)
+}
+
+func require_True(t *testing.T, b bool) {
+	t.Helper()
+	if !b {
+		t.Fatalf("require true, but got false")
+	}
+}
+
+func require_Equal[T comparable](t *testing.T, a, b T) {
+	t.Helper()
+	if a != b {
+		t.Fatalf("require %T equal, but got: %v != %v", a, a, b)
+	}
+}

server/thw/thw.go

@@ -0,0 +1,244 @@
+// Copyright 2024-2026 The NATS Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package thw
+
+import (
+	"encoding/binary"
+	"errors"
+	"io"
+	"math"
+	"time"
+)
+
+// Error for when we can not locate a task for removal or updates.
+var ErrTaskNotFound = errors.New("thw: task not found")
+
+// Error for when we try to decode a binary-encoded THW with an unknown version number.
+var ErrInvalidVersion = errors.New("thw: encoded version not known")
+
+const (
+	tickDuration = int64(time.Second) // Tick duration in nanoseconds.
+	wheelBits    = 12                 // 2^12 = 4096 slots.
+	wheelSize    = 1 << wheelBits     // Number of slots in the wheel.
+	wheelMask    = wheelSize - 1      // Mask for calculating position.
+	headerLen    = 17                 // 1 byte magic + 2x uint64s
+)
+
+// slot represents a single slot in the wheel.
+type slot struct {
+	entries map[uint64]int64 // Map of sequence to expires.
+	lowest  int64            // Lowest expiration time in this slot.
+}
+
+// HashWheel represents the timing wheel.
+type HashWheel struct {
+	wheel  []*slot // Array of slots.
+	lowest int64   // Track the lowest expiration time across all slots.
+	count  uint64  // How many entries are present?
+}
+
+// HashWheelEntry represents a single entry in the wheel.
+type HashWheelEntry struct {
+	Seq     uint64
+	Expires int64
+}
+
+// NewHashWheel initializes a new HashWheel.
+func NewHashWheel() *HashWheel {
+	return &HashWheel{
+		wheel:  make([]*slot, wheelSize),
+		lowest: math.MaxInt64,
+	}
+}
+
+// getPosition calculates the slot position for a given expiration time.
+func (hw *HashWheel) getPosition(expires int64) int64 {
+	return (expires / tickDuration) & wheelMask
+}
+
+// newSlot creates a new slot.
+func newSlot() *slot {
+	return &slot{
+		entries: make(map[uint64]int64),
+		lowest:  math.MaxInt64,
+	}
+}
+
+// Add schedules a new timer task.
+func (hw *HashWheel) Add(seq uint64, expires int64) error {
+	pos := hw.getPosition(expires)
+	// Initialize the slot lazily.
+	if hw.wheel[pos] == nil {
+		hw.wheel[pos] = newSlot()
+	}
+	if _, ok := hw.wheel[pos].entries[seq]; !ok {
+		hw.count++
+	}
+	hw.wheel[pos].entries[seq] = expires
+
+	// Update slot's lowest expiration if this is earlier.
+	if expires < hw.wheel[pos].lowest {
+		hw.wheel[pos].lowest = expires
+		// Update global lowest if this is now the earliest.
+		if expires < hw.lowest {
+			hw.lowest = expires
+		}
+	}
+
+	return nil
+}
+
+// Remove removes a timer task.
+func (hw *HashWheel) Remove(seq uint64, expires int64) error {
+	pos := hw.getPosition(expires)
+	s := hw.wheel[pos]
+	if s == nil {
+		return ErrTaskNotFound
+	}
+	if _, exists := s.entries[seq]; !exists {
+		return ErrTaskNotFound
+	}
+	delete(s.entries, seq)
+	hw.count--
+
+	// If the slot is empty, we can set it to nil to free memory.
+	if len(s.entries) == 0 {
+		hw.wheel[pos] = nil
+	}
+	return nil
+}
+
+// Update updates the expiration time of an existing timer task.
+func (hw *HashWheel) Update(seq uint64, oldExpires int64, newExpires int64) error {
+	// Remove from old position.
+	if err := hw.Remove(seq, oldExpires); err != nil {
+		return err
+	}
+	// Add to new position.
+	return hw.Add(seq, newExpires)
+}
+
+// ExpireTasks processes all expired tasks using a callback, but only expires a task if the callback returns true.
+func (hw *HashWheel) ExpireTasks(callback func(seq uint64, expires int64) bool) {
+	now := time.Now().UnixNano()
+	hw.expireTasks(now, callback)
+}
+
+func (hw *HashWheel) expireTasks(ts int64, callback func(seq uint64, expires int64) bool) {
+	// Quick return if nothing is expired.
+	if hw.lowest > ts {
+		return
+	}
+
+	globalLowest := int64(math.MaxInt64)
+	for pos, s := range hw.wheel {
+		// Skip s if nothing to expire.
+		if s == nil || s.lowest > ts {
+			if s != nil && s.lowest < globalLowest {
+				globalLowest = s.lowest
+			}
+			continue
+		}
+
+		// Track new lowest while processing expirations
+		slotLowest := int64(math.MaxInt64)
+		for seq, expires := range s.entries {
+			if expires <= ts && callback(seq, expires) {
+				// Only remove if not done so already by the callback.
+				if _, ok := s.entries[seq]; ok {
+					delete(s.entries, seq)
+					hw.count--
+				}
+				continue
+			}
+			if expires < slotLowest {
+				slotLowest = expires
+			}
+		}
+
+		// Nil out if we are empty.
+		if len(s.entries) == 0 {
+			hw.wheel[pos] = nil
+		} else {
+			s.lowest = slotLowest
+			if slotLowest < globalLowest {
+				globalLowest = slotLowest
+			}
+		}
+	}
+	hw.lowest = globalLowest
+}
+
+// GetNextExpiration returns the earliest expiration time before the given time.
+// Returns math.MaxInt64 if no expirations exist before the specified time.
+func (hw *HashWheel) GetNextExpiration(before int64) int64 {
+	if hw.lowest < before {
+		return hw.lowest
+	}
+	return math.MaxInt64
+}
+
+// Count returns the amount of tasks in the THW.
+func (hw *HashWheel) Count() uint64 {
+	return hw.count
+}
+
+// Encode writes out the contents of the THW into a binary snapshot
+// and returns it. The high seq number is included in the snapshot and will
+// be returned on decode.
+func (hw *HashWheel) Encode(highSeq uint64) []byte {
+	b := make([]byte, 0, headerLen+(hw.count*(2*binary.MaxVarintLen64)))
+	b = append(b, 1)                                  // Magic version
+	b = binary.LittleEndian.AppendUint64(b, hw.count) // Entry count
+	b = binary.LittleEndian.AppendUint64(b, highSeq)  // Stamp
+	for _, slot := range hw.wheel {
+		if slot == nil || slot.entries == nil {
+			continue
+		}
+		for v, ts := range slot.entries {
+			b = binary.AppendVarint(b, ts)
+			b = binary.AppendUvarint(b, v)
+		}
+	}
+	return b
+}
+
+// Decode snapshots a binary-encoded THW and replaces the contents of this
+// THW with them. Returns the high seq number from the snapshot.
+func (hw *HashWheel) Decode(b []byte) (uint64, error) {
+	if len(b) < headerLen {
+		return 0, io.ErrShortBuffer
+	}
+	if b[0] != 1 {
+		return 0, ErrInvalidVersion
+	}
+	hw.wheel = make([]*slot, wheelSize)
+	hw.lowest = math.MaxInt64
+	count := binary.LittleEndian.Uint64(b[1:])
+	stamp := binary.LittleEndian.Uint64(b[9:])
+	b = b[headerLen:]
+	for i := uint64(0); i < count; i++ {
+		ts, tn := binary.Varint(b)
+		if tn <= 0 {
+			return 0, io.ErrUnexpectedEOF
+		}
+		v, vn := binary.Uvarint(b[tn:])
+		if vn <= 0 {
+			return 0, io.ErrUnexpectedEOF
+		}
+		hw.Add(v, ts)
+		b = b[tn+vn:]
+	}
+	return stamp, nil
+}

server/thw/thw_test.go

@@ -0,0 +1,387 @@
+// Copyright 2024-2026 The NATS Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package thw
+
+import (
+	"encoding/binary"
+	"io"
+	"math"
+	"testing"
+	"time"
+)
+
+func TestHashWheelBasics(t *testing.T) {
+	hw := NewHashWheel()
+	now := time.Now()
+
+	// Add a sequence.
+	seq, expires := uint64(1), now.Add(5*time.Second).UnixNano()
+	require_NoError(t, hw.Add(seq, expires))
+	require_Equal(t, hw.count, 1)
+
+	// Try to remove non-existent sequence.
+	require_Error(t, hw.Remove(999, expires), ErrTaskNotFound)
+	require_Equal(t, hw.count, 1)
+
+	// Remove the sequence properly.
+	require_NoError(t, hw.Remove(seq, expires))
+	require_Equal(t, hw.count, 0)
+
+	// Verify it's gone.
+	require_Error(t, hw.Remove(seq, expires), ErrTaskNotFound)
+	require_Equal(t, hw.count, 0)
+}
+
+func TestHashWheelUpdate(t *testing.T) {
+	hw := NewHashWheel()
+	now := time.Now()
+	oldExpires := now.Add(5 * time.Second).UnixNano()
+	newExpires := now.Add(10 * time.Second).UnixNano()
+
+	// Add initial sequence.
+	require_NoError(t, hw.Add(1, oldExpires))
+	require_Equal(t, hw.count, 1)
+
+	// Update expiration.
+	require_NoError(t, hw.Update(1, oldExpires, newExpires))
+	require_Equal(t, hw.count, 1)
+
+	// Verify old expiration is gone.
+	require_Error(t, hw.Remove(1, oldExpires), ErrTaskNotFound)
+	require_Equal(t, hw.count, 1)
+
+	// Verify new expiration exists
+	require_NoError(t, hw.Remove(1, newExpires))
+	require_Equal(t, hw.count, 0)
+}
+
+func TestHashWheelExpiration(t *testing.T) {
+	hw := NewHashWheel()
+	now := time.Now()
+
+	// Add sequences with different expiration times.
+	seqs := map[uint64]int64{
+		1: now.Add(-1 * time.Second).UnixNano(), // Already expired
+		2: now.Add(1 * time.Second).UnixNano(),  // Expires soon
+		3: now.Add(10 * time.Second).UnixNano(), // Expires later
+		4: now.Add(60 * time.Second).UnixNano(), // Expires much later
+	}
+
+	for seq, expires := range seqs {
+		require_NoError(t, hw.Add(seq, expires))
+	}
+	require_Equal(t, hw.count, uint64(len(seqs)))
+
+	// Process expired tasks.
+	expired := make(map[uint64]bool)
+	hw.ExpireTasks(func(seq uint64, expires int64) bool {
+		expired[seq] = true
+		return true
+	})
+
+	// Verify only sequence 1 expired.
+	require_Equal(t, len(expired), 1)
+	require_True(t, expired[1])
+	require_Equal(t, hw.count, 3)
+}
+
+func TestHashWheelManualExpiration(t *testing.T) {
+	hw := NewHashWheel()
+	now := time.Now().UnixNano()
+
+	for seq := uint64(1); seq <= 4; seq++ {
+		require_NoError(t, hw.Add(seq, now))
+	}
+	require_Equal(t, hw.count, 4)
+
+	// Loop over expired multiple times, but without removing them.
+	expired := make(map[uint64]uint64)
+	for i := uint64(0); i <= 1; i++ {
+		hw.ExpireTasks(func(seq uint64, expires int64) bool {
+			expired[seq]++
+			return false
+		})
+
+		require_Equal(t, len(expired), 4)
+		require_Equal(t, expired[1], 1+i)
+		require_Equal(t, expired[2], 1+i)
+		require_Equal(t, expired[3], 1+i)
+		require_Equal(t, expired[4], 1+i)
+		require_Equal(t, hw.count, 4)
+	}
+
+	// Only remove even sequences.
+	for i := uint64(0); i <= 1; i++ {
+		hw.ExpireTasks(func(seq uint64, expires int64) bool {
+			expired[seq]++
+			return seq%2 == 0
+		})
+
+		// Verify even sequences are removed.
+		require_Equal(t, expired[1], 3+i)
+		require_Equal(t, expired[2], 3)
+		require_Equal(t, expired[3], 3+i)
+		require_Equal(t, expired[4], 3)
+		require_Equal(t, hw.count, 2)
+	}
+
+	// Manually remove last items.
+	require_NoError(t, hw.Remove(1, now))
+	require_NoError(t, hw.Remove(3, now))
+	require_Equal(t, hw.count, 0)
+}
+
+func TestHashWheelExpirationLargerThanWheel(t *testing.T) {
+	hw := NewHashWheel()
+
+	// Add sequences such that they can be expired immediately.
+	seqs := map[uint64]int64{
+		1: 0,
+		2: int64(time.Second),
+	}
+	for seq, expires := range seqs {
+		require_NoError(t, hw.Add(seq, expires))
+	}
+	require_Equal(t, hw.count, 2)
+
+	// Pick a timestamp such that the expiration needs to wrap around the whole wheel.
+	now := int64(time.Second) * wheelMask
+
+	// Process expired tasks.
+	expired := make(map[uint64]bool)
+	hw.expireTasks(now, func(seq uint64, expires int64) bool {
+		expired[seq] = true
+		return true
+	})
+
+	// Verify both sequences are expired.
+	require_Equal(t, len(expired), 2)
+	require_Equal(t, hw.count, 0)
+}
+
+func TestHashWheelNextExpiration(t *testing.T) {
+	hw := NewHashWheel()
+	now := time.Now()
+
+	// Add sequences with different expiration times.
+	seqs := map[uint64]int64{
+		1: now.Add(5 * time.Second).UnixNano(),
+		2: now.Add(3 * time.Second).UnixNano(), // Earliest
+		3: now.Add(10 * time.Second).UnixNano(),
+	}
+
+	for seq, expires := range seqs {
+		require_NoError(t, hw.Add(seq, expires))
+	}
+	require_Equal(t, hw.count, uint64(len(seqs)))
+
+	// Test GetNextExpiration.
+	nextExternalTick := now.Add(6 * time.Second).UnixNano()
+	// Should return sequence 2's expiration
+	require_Equal(t, hw.GetNextExpiration(nextExternalTick), seqs[2])
+
+	// Test with empty wheel.
+	empty := NewHashWheel()
+	require_Equal(t, empty.GetNextExpiration(now.Add(1*time.Second).UnixNano()), math.MaxInt64)
+}
+
+func TestHashWheelStress(t *testing.T) {
+	hw := NewHashWheel()
+	now := time.Now()
+
+	// Add many sequences.
+	numSequences := 100_000
+	for seq := 0; seq < numSequences; seq++ {
+		expires := now.Add(time.Duration(seq) * time.Second).UnixNano()
+		require_NoError(t, hw.Add(uint64(seq), expires))
+	}
+
+	// Update many sequences.
+	for seq := 0; seq < numSequences; seq += 2 { // Update every other sequence
+		oldExpires := now.Add(time.Duration(seq) * time.Second).UnixNano()
+		newExpires := now.Add(time.Duration(seq+numSequences) * time.Second).UnixNano()
+		require_NoError(t, hw.Update(uint64(seq), oldExpires, newExpires))
+	}
+
+	// Remove many sequences.
+	for seq := 1; seq < numSequences; seq += 2 { // Remove odd-numbered sequences
+		expires := now.Add(time.Duration(seq) * time.Second).UnixNano()
+		require_NoError(t, hw.Remove(uint64(seq), expires))
+	}
+}
+
+func TestHashWheelEncodeDecode(t *testing.T) {
+	hw := NewHashWheel()
+	now := time.Now()
+
+	// Add many sequences.
+	numSequences := 100_000
+	for seq := 0; seq < numSequences; seq++ {
+		expires := now.Add(time.Duration(seq) * time.Second).UnixNano()
+		require_NoError(t, hw.Add(uint64(seq), expires))
+	}
+
+	b := hw.Encode(12345)
+	require_True(t, len(b) > 17) // Bigger than just the header
+
+	nhw := NewHashWheel()
+	stamp, err := nhw.Decode(b)
+	require_NoError(t, err)
+	require_Equal(t, stamp, 12345)
+	require_Equal(t, hw.GetNextExpiration(math.MaxInt64), nhw.GetNextExpiration(math.MaxInt64))
+
+	for s, slot := range hw.wheel {
+		nslot := nhw.wheel[s]
+		require_Equal(t, slot.lowest, nslot.lowest)
+		require_Equal(t, len(slot.entries), len(nslot.entries))
+		for v, ts := range slot.entries {
+			nts, ok := nslot.entries[v]
+			require_True(t, ok)
+			require_Equal(t, ts, nts)
+		}
+	}
+}
+
+func TestHashWheelDecodeRejectsMalformed(t *testing.T) {
+	// Build a valid header that claims the given number of entries.
+	header := func(count uint64) []byte {
+		b := make([]byte, headerLen)
+		b[0] = 1                                    // Magic version
+		binary.LittleEndian.PutUint64(b[1:], count) // Entry count
+		binary.LittleEndian.PutUint64(b[9:], 0)     // High sequence stamp
+		return b
+	}
+
+	for _, test := range []struct {
+		title string
+		buf   []byte
+		err   error
+	}{
+		{title: "ShortHeader", buf: make([]byte, headerLen-1), err: io.ErrShortBuffer},
+		{title: "BadVersion", buf: func() []byte { b := header(0); b[0] = 2; return b }(), err: ErrInvalidVersion},
+		// Claims one entry but the buffer ends right after the header.
+		{title: "TruncatedAtVarints", buf: header(1), err: io.ErrUnexpectedEOF},
+		// Has the timestamp varint but the seq varint is missing.
+		{title: "TruncatedSeq", buf: append(header(1), binary.AppendVarint(nil, 12345)...), err: io.ErrUnexpectedEOF},
+	} {
+		t.Run(test.title, func(t *testing.T) {
+			hw := NewHashWheel()
+			_, err := hw.Decode(test.buf)
+			require_Error(t, err, test.err)
+		})
+	}
+}
+
+// Benchmarks
+
+func BenchmarkHashWheel_Add(b *testing.B) {
+	hw := NewHashWheel()
+	now := time.Now()
+
+	// Create different ranges for expires to spread across slots
+	expires := make([]int64, b.N)
+	for i := 0; i < b.N; i++ {
+		expires[i] = now.Add(time.Duration(i%3600) * time.Second).UnixNano()
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		hw.Add(uint64(i), expires[i])
+	}
+}
+
+func BenchmarkHashWheel_Add_SameSlot(b *testing.B) {
+	hw := NewHashWheel()
+	expires := time.Now().Add(time.Second).UnixNano()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		hw.Add(uint64(i), expires)
+	}
+}
+
+func BenchmarkHashWheel_Update(b *testing.B) {
+	hw := NewHashWheel()
+	now := time.Now()
+
+	// First add N items
+	expires := make([]int64, b.N)
+	for i := 0; i < b.N; i++ {
+		expires[i] = now.Add(time.Duration(i%3600) * time.Second).UnixNano()
+		hw.Add(uint64(i), expires[i])
+	}
+
+	newExpires := now.Add(2 * time.Hour).UnixNano()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		hw.Update(uint64(i), expires[i], newExpires)
+	}
+}
+
+func BenchmarkHashWheel_Update_SameSlot(b *testing.B) {
+	hw := NewHashWheel()
+	now := time.Now()
+	oldExpires := now.Add(time.Second).UnixNano()
+	newExpires := now.Add(2 * time.Second).UnixNano()
+
+	// First add N items
+	for i := 0; i < b.N; i++ {
+		hw.Add(uint64(i), oldExpires)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		hw.Update(uint64(i), oldExpires, newExpires)
+	}
+}
+
+// Benchmark memory allocation
+func BenchmarkHashWheel_Add_Memory(b *testing.B) {
+	b.ReportAllocs()
+	hw := NewHashWheel()
+	now := time.Now()
+
+	expires := make([]int64, b.N)
+	for i := 0; i < b.N; i++ {
+		expires[i] = now.Add(time.Duration(i%3600) * time.Second).UnixNano()
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		hw.Add(uint64(i), expires[i])
+	}
+}
+
+// Large scale test
+func BenchmarkHashWheel_LargeScale(b *testing.B) {
+	hw := NewHashWheel()
+	now := time.Now()
+
+	// Pre-populate with 100k items
+	for i := 0; i < 100_000; i++ {
+		expires := now.Add(time.Duration(i%3600) * time.Second).UnixNano()
+		hw.Add(uint64(i), expires)
+	}
+
+	expires := make([]int64, b.N)
+	for i := 0; i < b.N; i++ {
+		expires[i] = now.Add(time.Duration(i%3600) * time.Second).UnixNano()
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		hw.Add(uint64(i+100000), expires[i])
+	}
+}