Use generational stack for idle connections

internal/genstack/gen_stack.go

@@ -0,0 +1,85 @@
+package genstack
+
+// GenStack implements a generational stack.
+//
+// GenStack works as common stack except for the fact that all elements in the
+// older generation are guaranteed to be popped before any element in the newer
+// generation. New elements are always pushed to the current (newest)
+// generation.
+//
+// We could also say that GenStack behaves as a stack in case of a single
+// generation, but it behaves as a queue of individual generation stacks.
+type GenStack[T any] struct {
+	// We can represent arbitrary number of generations using 2 stacks. The
+	// new stack stores all new pushes and the old stack serves all reads.
+	// Old stack can represent multiple generations. If old == new, then all
+	// elements pushed in previous (not current) generations have already
+	// been popped.
+
+	old *stack[T]
+	new *stack[T]
+}
+
+// NewGenStack creates a new empty GenStack.
+func NewGenStack[T any]() *GenStack[T] {
+	s := &stack[T]{}
+	return &GenStack[T]{
+		old: s,
+		new: s,
+	}
+}
+
+func (s *GenStack[T]) Pop() (T, bool) {
+	// Pushes always append to the new stack, so if the old once becomes
+	// empty, it will remail empty forever.
+	if s.old.len() == 0 && s.old != s.new {
+		s.old = s.new
+	}
+
+	if s.old.len() == 0 {
+		var zero T
+		return zero, false
+	}
+
+	return s.old.pop(), true
+}
+
+// Push pushes a new element at the top of the stack.
+func (s *GenStack[T]) Push(v T) { s.new.push(v) }
+
+// NextGen starts a new stack generation.
+func (s *GenStack[T]) NextGen() {
+	if s.old == s.new {
+		s.new = &stack[T]{}
+		return
+	}
+
+	// We need to pop from the old stack to the top of the new stack. Let's
+	// have an example:
+	//
+	//   Old: <bottom> 4 3 2 1
+	//   New: <bottom> 8 7 6 5
+	//   PopOrder: 1 2 3 4 5 6 7 8
+	//
+	//
+	// To preserve pop order, we have to take all elements from the old
+	// stack and push them to the top of new stack:
+	//
+	//   New: 8 7 6 5 4 3 2 1
+	//
+	s.new.push(s.old.takeAll()...)
+
+	// We have the old stack allocated and empty, so why not to reuse it as
+	// new new stack.
+	s.old, s.new = s.new, s.old
+}
+
+// Len returns number of elements in the stack.
+func (s *GenStack[T]) Len() int {
+	l := s.old.len()
+	if s.old != s.new {
+		l += s.new.len()
+	}
+
+	return l
+}

internal/genstack/gen_stack_test.go

@@ -0,0 +1,90 @@
+package genstack
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func requirePopEmpty[T any](t testing.TB, s *GenStack[T]) {
+	v, ok := s.Pop()
+	require.False(t, ok)
+	require.Zero(t, v)
+}
+
+func requirePop[T any](t testing.TB, s *GenStack[T], expected T) {
+	v, ok := s.Pop()
+	require.True(t, ok)
+	require.Equal(t, expected, v)
+}
+
+func TestGenStack_Empty(t *testing.T) {
+	s := NewGenStack[int]()
+	requirePopEmpty(t, s)
+}
+
+func TestGenStack_SingleGen(t *testing.T) {
+	r := require.New(t)
+	s := NewGenStack[int]()
+
+	s.Push(1)
+	s.Push(2)
+	r.Equal(2, s.Len())
+
+	requirePop(t, s, 2)
+	requirePop(t, s, 1)
+	requirePopEmpty(t, s)
+}
+
+func TestGenStack_TwoGen(t *testing.T) {
+	r := require.New(t)
+	s := NewGenStack[int]()
+
+	s.Push(3)
+	s.Push(4)
+	s.Push(5)
+	r.Equal(3, s.Len())
+	s.NextGen()
+	r.Equal(3, s.Len())
+	s.Push(6)
+	s.Push(7)
+	r.Equal(5, s.Len())
+
+	requirePop(t, s, 5)
+	requirePop(t, s, 4)
+	requirePop(t, s, 3)
+	requirePop(t, s, 7)
+	requirePop(t, s, 6)
+	requirePopEmpty(t, s)
+}
+
+func TestGenStack_MuptiGen(t *testing.T) {
+	r := require.New(t)
+	s := NewGenStack[int]()
+
+	s.Push(10)
+	s.Push(11)
+	s.Push(12)
+	r.Equal(3, s.Len())
+	s.NextGen()
+	r.Equal(3, s.Len())
+	s.Push(13)
+	s.Push(14)
+	r.Equal(5, s.Len())
+	s.NextGen()
+	r.Equal(5, s.Len())
+	s.Push(15)
+	s.Push(16)
+	s.Push(17)
+	r.Equal(8, s.Len())
+
+	requirePop(t, s, 12)
+	requirePop(t, s, 11)
+	requirePop(t, s, 10)
+	requirePop(t, s, 14)
+	requirePop(t, s, 13)
+	requirePop(t, s, 17)
+	requirePop(t, s, 16)
+	requirePop(t, s, 15)
+	requirePopEmpty(t, s)
+}

internal/genstack/stack.go

@@ -0,0 +1,39 @@
+package genstack
+
+// stack is a wrapper around an array implementing a stack.
+//
+// We cannot use slice to represent the stack because append might change the
+// pointer value of the slice. That would be an issue in GenStack
+// implementation.
+type stack[T any] struct {
+	arr []T
+}
+
+// push pushes a new element at the top of a stack.
+func (s *stack[T]) push(vs ...T) { s.arr = append(s.arr, vs...) }
+
+// pop pops the stack top-most element.
+//
+// If stack length is zero, this method panics.
+func (s *stack[T]) pop() T {
+	idx := s.len() - 1
+	val := s.arr[idx]
+
+	// Avoid memory leak
+	var zero T
+	s.arr[idx] = zero
+
+	s.arr = s.arr[:idx]
+	return val
+}
+
+// takeAll returns all elements in the stack in order as they are stored - i.e.
+// the top-most stack element is the last one.
+func (s *stack[T]) takeAll() []T {
+	arr := s.arr
+	s.arr = nil
+	return arr
+}
+
+// len returns number of elements in the stack.
+func (s *stack[T]) len() int { return len(s.arr) }