support keyed pool items

adaptive_test.go

@@ -54,7 +54,7 @@ func TestAdaptive(t *testing.T) {
 		var o pool.Options
 		o.Testing.Clock = clock
 		o.Testing.Bus = bus
-		alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+		alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 		p := pool.Make(alloc, nil, o)
 
 		var active [][]byte
@@ -127,7 +127,7 @@ func TestAdaptive(t *testing.T) {
 				steadyUseCycles = 8
 			)
 
-			alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+			alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 			p := pool.Make(alloc, nil, pool.Options{})
 
 			var active [][]byte
@@ -174,7 +174,7 @@ func TestAdaptive(t *testing.T) {
 				highLoadCycles  = 8
 			)
 
-			alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+			alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 			p := pool.Make(alloc, nil, pool.Options{})
 
 			var active [][]byte
@@ -231,7 +231,7 @@ func TestAdaptive(t *testing.T) {
 				steadyUseCycles = 8
 			)
 
-			alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+			alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 			p := pool.Make(alloc, nil, pool.Options{})
 
 			var active [][]byte
@@ -280,7 +280,7 @@ func TestAdaptive(t *testing.T) {
 				highLoadCycles  = 8
 			)
 
-			alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+			alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 			p := pool.Make(alloc, nil, pool.Options{})
 
 			var active [][]byte

event_test.go

@@ -8,7 +8,7 @@ import (
 
 func TestEvent(t *testing.T) {
 	t.Run("get put allocate free", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+		alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 		e := make(chan pool.Event, 7)
 		o := pool.Options{
 			Algo:      pool.Max(2),
@@ -53,7 +53,7 @@ func TestEvent(t *testing.T) {
 	})
 
 	t.Run("allocate error", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return nil, errTest }
+		alloc := func(string) ([]byte, error) { return nil, errTest }
 		e := make(chan pool.Event, 1)
 		o := pool.Options{
 			Algo:      pool.NoShrink(),
@@ -74,7 +74,7 @@ func TestEvent(t *testing.T) {
 	})
 
 	t.Run("drop", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+		alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 		e := make(chan pool.Event)
 		o := pool.Options{
 			Algo:      pool.NoShrink(),
@@ -96,7 +96,7 @@ func TestEvent(t *testing.T) {
 	})
 
 	t.Run("mask", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+		alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 		e := make(chan pool.Event, 1)
 		o := pool.Options{
 			Algo:      pool.NoShrink(),

lib.go

@@ -73,7 +73,7 @@ const (
 	FreeOperation
 
 	// AllocateError is the type of events sent after a failed allocation. The error will not be included
-	// with the event, but it will be returned by the failed Get function call.
+	// with the event, but it will be returned by the failed Get* function call.
 	AllocateError
 
 	// AllEvents can be used as a mask that includes all the event types.
@@ -101,7 +101,7 @@ type Event struct {
 // the pool guarantees that the algo instance is accessed only from one goroutine at a time.
 type Algo interface {
 
-	// Target is called on every Put operation with the current pool state as the input. It is expected to
+	// Target is called on every Put* operation with the current pool state as the input. It is expected to
 	// return the target idle count, as dictated by the implementing algorithm. Optionally, a timeout value
 	// can be returned (nextCheck), and if it is a positive value, the pool will call Target again after the
 	// defined time expires to fetch the next target idle count. In each case, when Target returns a smaller
@@ -110,8 +110,8 @@ type Algo interface {
 	// When using nextCheck, not every returned value results in calling Target by the pool, only the ones
 	// that were set after the previous check expired.
 	//
-	// Implementations should consider that while the recommended way of using the pool is to only call Put
+	// Implementations should consider that while the recommended way of using the pool is to only call Put*
-	// with items that were received by calling Get, the pool itself doesn't prohibit calling Put with
+	// with items that were received by calling Get*, the pool itself doesn't prohibit calling Put* with
 	// 'foreign' items.
 	Target(Stats) (target int, nextCheck time.Duration)
 
@@ -141,12 +141,12 @@ type Options struct {
 
 	Testing struct {
 
-		// Clock is an optional clock meant to be used with testing. The main purpose is to avoid time sensitive
+		// Clock is an optional clock meant to be used with testing. The main purpose is to avoid time
-		// tests running for a too long time.
+		// sensitive tests running for a too long time.
 		Clock times.Clock
 
-		// Bus is an optional signal bus meant to be used with testing. The main purpose is to ensure that
+		// Bus is an optional signal bus meant to be used with testing. The main purpose is to ensure
-		// specific blocks of code are executed in a predefined order during concurrent tests.
+		// that specific blocks of code are executed in a predefined order during concurrent tests.
 		Bus *syncbus.SyncBus
 	}
 }
@@ -158,8 +158,8 @@ type Pool[R any] struct {
 	pool pool[R]
 }
 
-// ErrNoItems is returned on Get calls when the pool does not have enough items and it was initialized without an
+// ErrNoItems is returned on Get* calls when the pool does not have enough items and it was initialized without
-// allocate function.
+// an allocate function.
 var ErrNoItems = errors.New("not enough items in the pool")
 
 // ErrUninitialized is returned when the pool was not initialized.
@@ -230,7 +230,7 @@ func (s Stats) String() string {
 // on very sudden drops in traffic, it ensures the eventual release of all pooled items with an internal
 // background job. This job is timed based on the duration of the last active usage session (the period during
 // which there were active items). Together with the pool implementation, it always reuses the most recent
-// items, as in LIFO for Get and FIFO for Free.
+// items, as in LIFO for Get* and FIFO for Free.
 //
 // We need to be aware of some potential caveats due to its zero-config nature. Because it relies on the
 // sequence of operations rather than wall-clock thresholds for its core logic, the algorithm treats rapid
@@ -243,7 +243,7 @@ func Adaptive() Algo {
 
 // MaxTimeout creates a pool shrink algorithm instance, that releases items whenever the number of idle items
 // exceeds max, and it also releases those items that were idle for too long. Together with the pool, it ensures
-// that the Get operation is LIFO and the Free operation is FIFO.
+// that the Get* operation is LIFO and the Free operation is FIFO.
 //
 // If max <= 0, the max pool size is not enforced. If to <= 0, the timeout is not enforced.
 func MaxTimeout(max int, to time.Duration) Algo {
@@ -268,23 +268,55 @@ func NoShrink() Algo {
 
 // Make initializes a Pool instance.
 //
-// The parameter alloc is used on Get operations when the pool is empty. If alloc is nil, and the pool is empty
+// The parameter alloc is used on Get* operations when the pool is empty. For unkeyed Get* operations, the input
-// at the time of calling Get, Get will return ErrEmpty. If alloc returns an error, the same error is returned
+// key will be the empty string. If alloc is nil, and the pool is empty at the time of calling Get*, Get* will
-// by Get. If events were configured, alloc triggers an AllocateOperation event. This event is typically the
+// return ErrEmpty. If alloc returns an error, the same error is returned by Get*. If events were configured,
-// same as the GetOperation event.
+// alloc triggers an AllocateOperation event. This event is typically the same as the GetOperation event.
 //
 // The parameter free is called when an item is released from the pool, with the item being released as the
 // argument. It can be nil for resource types that don't need explicit deallocation. If events were configured,
 // releasing an item triggers a FreeOperation event, regardless if the free parameter is nil.
-func Make[R any](alloc func() (R, error), free func(R), o Options) Pool[R] {
+func Make[R any](alloc func(string) (R, error), free func(R), o Options) Pool[R] {
 	return Pool[R]{pool: makePool(alloc, free, o)}
 }
 
-// GetN returns n items from the pool. If the pool is empty and no allocation function was configured, it
+// GetKeyN returns n items with the provided key from the pool. If the pool has no items with the provided key
-// returns ErrEmpty. If the pool has less than n items and no allocation function was configured, it returns the
+// and no allocation function was configured, it returns ErrEmpty. If the pool has less than n items with the
-// available items and ErrEtmpy. If the pool is empty or has less than the requested items, and the allocation
+// provided key and no allocation function was configured, it returns ErrEmpty. If the pool has no items with
-// function returns an error, it returns the available items and that error. If events were configured, GetN
+// the provided key or has less than the requested items, and the allocation function returns an error, it
+// returns the available items and that error. If events were configured, GetKeyN triggers a GetOperation event.
+func (p Pool[R]) GetKeyN(k string, n int) ([]R, error) {
+	if p.pool.state == nil {
+		return nil, ErrUninitialized
+	}
+
+	return p.pool.getKey(key{value: k}, n)
+}
+
+// GetKey returns an item from the pool with the provided key. If the pool has no items with the provided key
+// and no allocation function was configured, it returns ErrEmpty. If the pool has no items with the provided
+// key, and the allocation function returns an error, it returns that error. If events were configured, GetKey
 // triggers a GetOperation event.
+func (p Pool[R]) GetKey(k string) (R, error) {
+	if p.pool.state == nil {
+		var r R
+		return r, ErrUninitialized
+	}
+
+	r, err := p.pool.getKey(key{value: k}, 1)
+	if err != nil {
+		var r R
+		return r, err
+	}
+
+	return r[0], nil
+}
+
+// GetN returns n items from the pool. If the pool is empty and no allocation function was configured, it
+// returns ErrEmpty. If the pool has less than n items and no allocation function was configured, it returns
+// ErrEmpty. If the pool is empty or has less than the requested items, and the allocation function returns an
+// error, it returns the available items and that error. If events were configured, GetN triggers a GetOperation
+// event.
 func (p Pool[R]) GetN(n int) ([]R, error) {
 	if p.pool.state == nil {
 		return nil, ErrUninitialized
@@ -304,13 +336,28 @@ func (p Pool[R]) Get() (R, error) {
 
 	r, err := p.pool.get(1)
 	if err != nil {
-		var rr R
+		var r R
-		return rr, err
+		return r, err
 	}
 
 	return r[0], nil
 }
 
+// PutKey stores one or more items in the pool with the provided key. If events were configured, it triggers a
+// PutOperation event.
+//
+// It is recommended to use it only with items that were received by the GetKey method. While it is allowed to
+// put other items in the pool, it may change the way the shrinking algorithm works. E.g. it can be considered
+// as a sudden drop in the number of active items. If the pool needs to be prewarmed, or prepared for an
+// expected spike of traffic, consider using the Load method.
+func (p Pool[R]) PutKey(k string, i ...R) {
+	if p.pool.state == nil {
+		return
+	}
+
+	p.pool.putKey(key{value: k}, i...)
+}
+
 // Put stores one or more items in the pool. If events were configured, it triggers a PutOperation event.
 //
 // It is recommended to use it only with items that were received by the Get method. While it is allowed to put
@@ -325,7 +372,8 @@ func (p Pool[R]) Put(i ...R) {
 	p.pool.put(i...)
 }
 
-// DropN indicates to the pool that N items received by Get were dropped and will not be put back into the pool.
+// DropN indicates to the pool that N items received by Get* were dropped and will not be put back into the
+// pool.
 //
 // While it's not mandatory call it when some items were dropped, it has a twofold purpose. Depending on the
 // shrinking algorithm, it may work more consistently, if the algorithm is notified about the change in the
@@ -347,6 +395,17 @@ func (p Pool[R]) Drop() {
 	p.pool.drop(1)
 }
 
+// LoadKey can be used to populate the pool with items for the provided key that were not allocated as the
+// result of the GetKey operation. It can be useful in scenarios where prewarming or preparing for a sudden
+// traffic spike is necessary. If events were configured, it triggers a LoadOperation event.
+func (p Pool[R]) LoadKey(k string, i []R) {
+	if p.pool.state == nil {
+		return
+	}
+
+	p.pool.loadKey(key{value: k}, i)
+}
+
 // Load can be used to populate the pool with items that were not allocated as the result of the Get operation.
 // It can be useful in scenarios where prewarming or preparing for a sudden traffic spike is necessary. If
 // events were configured, it triggers a LoadOperation event.

maxto_test.go

@@ -225,7 +225,7 @@ func TestMaxTO(t *testing.T) {
 			var o pool.Options
 			o.Algo = pool.MaxTimeout(15, 300*time.Millisecond)
 			o.Testing.Clock = clock
-			alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+			alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 			p := pool.Make(alloc, nil, o)
 
 			var active [][]byte
@@ -280,7 +280,7 @@ func TestMaxTO(t *testing.T) {
 			var o pool.Options
 			o.Algo = pool.MaxTimeout(15, 300*time.Millisecond)
 			o.Testing.Clock = clock
-			alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+			alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 			p := pool.Make(alloc, nil, o)
 
 			var active [][]byte
@@ -348,7 +348,7 @@ func TestMaxTO(t *testing.T) {
 			var o pool.Options
 			o.Algo = pool.MaxTimeout(15, 300*time.Millisecond)
 			o.Testing.Clock = clock
-			alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+			alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 			p := pool.Make(alloc, nil, o)
 
 			var active [][]byte
@@ -405,7 +405,7 @@ func TestMaxTO(t *testing.T) {
 			var o pool.Options
 			o.Algo = pool.MaxTimeout(15, 300*time.Millisecond)
 			o.Testing.Clock = clock
-			alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+			alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 			p := pool.Make(alloc, nil, o)
 
 			var active [][]byte

pool.go

@@ -2,23 +2,33 @@ package pool
 
 import (
 	"code.squareroundforest.org/arpio/times"
+	"errors"
 	"time"
 )
 
+type key struct {
+	null  bool
+	value string
+}
+
 type state[R any] struct {
-	items              []R
+	keys               []key
+	items              map[key][]R
 	stats              Stats
 	forcedCheckPending bool
+	deletedKeys        int
 }
 
 type pool[R any] struct {
 	state   chan state[R]
-	alloc   func() (R, error)
+	alloc   func(string) (R, error)
 	free    func(r R)
 	options Options
 }
 
-func makePool[R any](alloc func() (R, error), free func(r R), o Options) pool[R] {
+var nullKey = key{null: true}
+
+func makePool[R any](alloc func(string) (R, error), free func(r R), o Options) pool[R] {
 	if o.Algo == nil {
 		o.Algo = Adaptive()
 	}
@@ -32,7 +42,7 @@ func makePool[R any](alloc func() (R, error), free func(r R), o Options) pool[R]
 	}
 
 	s := make(chan state[R], 1)
-	s <- state[R]{}
+	s <- state[R]{items: make(map[key][]R)}
 	return pool[R]{
 		state:   s,
 		alloc:   alloc,
@@ -70,7 +80,68 @@ func (p pool[R]) sendEvent(e EventType, s Stats) {
 	}
 }
 
-func (p pool[R]) get(n int) ([]R, error) {
+func appendKeys(all []key, k key, n int) []key {
+	for i := 0; i < n; i++ {
+		all = append(all, k)
+	}
+
+	return all
+}
+
+func dropKeys(all []key, k key, n int) []key {
+	first := len(all)
+	last := first
+	for n > 0 {
+		if first == 0 && last == 0 {
+			panic(errors.New("damaged pool: missing keys"))
+		}
+
+		if first > 0 && all[first-1] == k && last-first < n {
+			first--
+			continue
+		}
+
+		if first == last {
+			first--
+			last--
+			continue
+		}
+
+		var zero key
+		copy(all[first:], all[last:])
+		for i := len(all) + first - last; i < len(all); i++ {
+			all[i] = zero
+		}
+
+		all = all[:len(all)+first-last]
+		n += first - last
+		first--
+		last = first
+	}
+
+	return all
+}
+
+func resetItemsMap[R any](s state[R]) state[R] {
+	if s.deletedKeys < 1<<12 {
+		return s
+	}
+
+	if s.deletedKeys <= len(s.items)*2 {
+		return s
+	}
+
+	next := make(map[key][]R)
+	for k := range s.items {
+		next[k] = s.items[k]
+	}
+
+	s.items = next
+	s.deletedKeys = 0
+	return s
+}
+
+func (p pool[R]) getKey(k key, n int) ([]R, error) {
 	if n <= 0 {
 		return nil, nil
 	}
@@ -83,63 +154,76 @@ func (p pool[R]) get(n int) ([]R, error) {
 	var event EventType
 	event |= GetOperation
 	s.stats.Get++
-	if len(s.items) < n && p.alloc == nil {
+	if len(s.items[k]) < n && p.alloc == nil {
 		return nil, ErrNoItems
 	}
 
 	r := make([]R, n)
-	if len(s.items) <= n {
+	if len(s.items[k]) <= n {
-		copy(r, s.items)
+		copy(r, s.items[k])
 	}
 
-	if len(s.items) > n {
+	if len(s.items[k]) > n {
 		// using the items from the end of the list:
-		copy(r, s.items[len(s.items)-n:])
+		copy(r, s.items[k][len(s.items[k])-n:])
 	}
 
-	if len(s.items) < n {
+	if len(s.items[k]) < n {
 		var err error
 		event |= AllocateOperation
-		for i := len(s.items); i < n; i++ {
+		for i := len(s.items[k]); i < n; i++ {
 			s.stats.Alloc++
-			r[i], err = p.alloc()
+			if r[i], err = p.alloc(k.value); err == nil {
-			if err == nil {
 				continue
 			}
 
-			s.items = append(s.items, r[len(s.items):i]...)
+			s.items[k] = append(s.items[k], r[len(s.items[k]):i]...)
-			s.stats.Idle = len(s.items)
+			s.keys = appendKeys(s.keys, k, i-len(s.items[k]))
+			s.stats.Idle = len(s.keys)
 			event |= AllocateError
 			p.sendEvent(event, s.stats)
 			return nil, err
 		}
 	}
 
-	if len(s.items) <= n {
+	// drop enough keys from the end:
-		var zero R
+	d := n
-		for i := 0; i < len(s.items); i++ {
+	if len(s.items[k]) < n {
-			s.items[i] = zero
+		d = len(s.items[k])
-		}
-
-		s.items = nil
 	}
 
-	if len(s.items) > n {
+	s.keys = dropKeys(s.keys, k, d)
+	if len(s.items[k]) <= n {
 		var zero R
-		for i := len(s.items) - n; i < len(s.items); i++ {
+		for i := 0; i < len(s.items[k]); i++ {
-			s.items[i] = zero
+			s.items[k][i] = zero
 		}
 
-		s.items = s.items[:len(s.items)-n]
+		delete(s.items, k)
+		s.deletedKeys++
 	}
 
-	s.stats.Idle = len(s.items)
+	s = resetItemsMap(s)
+	if len(s.items[k]) > n {
+		var zero R
+		for i := len(s.items[k]) - n; i < len(s.items[k]); i++ {
+			s.items[k][i] = zero
+		}
+
+		s.items[k] = s.items[k][:len(s.items[k])-n]
+	}
+
+	s.stats.Idle = len(s.keys)
 	s.stats.Active += n
 	p.sendEvent(event, s.stats)
 	return r, nil
 }
 
-func (p pool[R]) put(r ...R) {
+func (p pool[R]) get(n int) ([]R, error) {
+	return p.getKey(nullKey, n)
+}
+
+func (p pool[R]) putKey(k key, r ...R) {
 	if len(r) == 0 {
 		return
 	}
@@ -162,25 +246,31 @@ func (p pool[R]) put(r ...R) {
 
 	t, f := p.options.Algo.Target(s.stats)
 	switch {
-	case t >= len(s.items)+len(r):
+	case t >= len(s.keys)+len(r):
-		s.items = append(s.items, r...)
+		s.items[k] = append(s.items[k], r...)
-	case t > len(s.items) && t < len(s.items)+len(r):
+		s.keys = appendKeys(s.keys, k, len(r))
+	case t > len(s.keys) && t < len(s.keys)+len(r):
 		event |= FreeOperation
-		s = p.freeItems(s, r[t-len(s.items):])
+		s = p.freeItems(s, r[t-len(s.keys):])
-		s.items = append(s.items, r[:t-len(s.items)]...)
+		s.items[k] = append(s.items[k], r[:t-len(s.keys)]...)
+		s.keys = appendKeys(s.keys, k, t-len(s.keys))
 	default:
 		event |= FreeOperation
 		s = p.freeItems(s, r)
 		s = p.freeAboveTarget(s, t)
 	}
 
-	s.stats.Idle = len(s.items)
+	s.stats.Idle = len(s.keys)
 	p.sendEvent(event, s.stats)
 	if f > 0 {
 		s = p.forcedCheck(s, f)
 	}
 }
 
+func (p pool[R]) put(r ...R) {
+	p.putKey(nullKey, r...)
+}
+
 func (p pool[R]) drop(n int) {
 	if n <= 0 {
 		return
@@ -200,20 +290,20 @@ func (p pool[R]) drop(n int) {
 	}
 
 	t, f := p.options.Algo.Target(s.stats)
-	if t < len(s.items) {
+	if t < len(s.keys) {
 		event |= FreeOperation
 		s = p.freeAboveTarget(s, t)
 	}
 
-	s.stats.Idle = len(s.items)
+	s.stats.Idle = len(s.keys)
 	p.sendEvent(event, s.stats)
 	if f > 0 {
 		s = p.forcedCheck(s, f)
 	}
 }
 
-func (p pool[R]) load(i []R) {
+func (p pool[R]) loadKey(k key, r []R) {
-	if len(i) == 0 {
+	if len(r) == 0 {
 		return
 	}
 
@@ -222,14 +312,19 @@ func (p pool[R]) load(i []R) {
 		p.state <- s
 	}()
 
-	s.items = append(s.items, i...)
+	s.items[k] = append(s.items[k], r...)
-	s.stats.Idle = len(s.items)
+	s.keys = appendKeys(s.keys, k, len(r))
-	s.stats.Load += len(i)
+	s.stats.Idle = len(s.keys)
-	p.options.Algo.Load(len(i))
+	s.stats.Load += len(r)
+	p.options.Algo.Load(len(r))
 	event := LoadOperation
 	p.sendEvent(event, s.stats)
 }
 
+func (p pool[R]) load(i []R) {
+	p.loadKey(nullKey, i)
+}
+
 func (p pool[R]) forcedCheck(s state[R], timeout time.Duration) state[R] {
 	if s.forcedCheckPending {
 		return s
@@ -260,23 +355,38 @@ func (p pool[R]) freeItems(s state[R], r []R) state[R] {
 }
 
 func (p pool[R]) freeAboveTarget(s state[R], target int) state[R] {
-	if len(s.items) <= target {
+	if len(s.keys) <= target {
 		return s
 	}
 
-	f := s.items[:len(s.items)-target]
+	var (
+		zeroKey  key
+		zeroItem R
+	)
 
-	var zero R
+	d := len(s.keys) - target
-	for i := 0; i < len(s.items)-target; i++ {
+	r := make([]R, 0, d)
-		s.items[i] = zero
+	for _, k := range s.keys[:d] {
+		if len(s.items[k]) == 0 {
+			panic(errors.New("damaged pool: missing item"))
+		}
+
+		r = append(r, s.items[k][0])
+		s.items[k][0] = zeroItem
+		s.items[k] = s.items[k][1:]
+		if len(s.items[k]) == 0 {
+			delete(s.items, k)
+			s.deletedKeys++
+		}
 	}
 
-	s.items = s.items[len(s.items)-target:]
+	s = resetItemsMap(s)
-	if len(s.items) == 0 {
+	for i := 0; i < d; i++ {
-		s.items = nil
+		s.keys[i] = zeroKey
 	}
 
-	return p.freeItems(s, f)
+	s.keys = s.keys[d:]
+	return p.freeItems(s, r)
 }
 
 func (p pool[R]) checkAndFree() {
@@ -289,7 +399,7 @@ func (p pool[R]) checkAndFree() {
 	t, f := p.options.Algo.Target(s.stats)
 	prev := s.stats.Free
 	s = p.freeAboveTarget(s, t)
-	s.stats.Idle = len(s.items)
+	s.stats.Idle = len(s.keys)
 	if s.stats.Free > prev {
 		p.sendEvent(FreeOperation, s.stats)
 	}
@@ -309,8 +419,13 @@ func (p pool[R]) freePool() {
 
 	s.stats.Idle = 0
 	prev := s.stats.Free
-	s = p.freeItems(s, s.items)
+	for k := range s.items {
-	s.items = nil
+		s = p.freeItems(s, s.items[k])
+	}
+
+	s.keys = nil
+	s.items = make(map[key][]R)
+	s.deletedKeys = 0
 	if s.stats.Free > prev {
 		p.sendEvent(FreeOperation, s.stats)
 	}

pool_test.go

@@ -13,7 +13,7 @@ var errTest = errors.New("test error")
 
 func TestPool(t *testing.T) {
 	t.Run("initial stats", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+		alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 		p := pool.Make(alloc, nil, pool.Options{Algo: pool.NoShrink()})
 		s := p.Stats()
 		e := pool.Stats{}
@@ -23,7 +23,7 @@ func TestPool(t *testing.T) {
 	})
 
 	t.Run("get when empty", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+		alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 		p := pool.Make(alloc, nil, pool.Options{Algo: pool.NoShrink()})
 		b, err := p.Get()
 		if err != nil {
@@ -42,7 +42,7 @@ func TestPool(t *testing.T) {
 	})
 
 	t.Run("get pooled own", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+		alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 		p := pool.Make(alloc, nil, pool.Options{Algo: pool.NoShrink()})
 		b, err := p.Get()
 		if err != nil {
@@ -71,7 +71,7 @@ func TestPool(t *testing.T) {
 	})
 
 	t.Run("get pooled foreign", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+		alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 		p := pool.Make(alloc, nil, pool.Options{Algo: pool.NoShrink()})
 		p.Put(make([]byte, 1<<6))
 		b, err := p.Get()
@@ -124,7 +124,7 @@ func TestPool(t *testing.T) {
 	})
 
 	t.Run("allocation error", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return nil, errTest }
+		alloc := func(string) ([]byte, error) { return nil, errTest }
 		p := pool.Make(alloc, nil, pool.Options{Algo: pool.NoShrink()})
 		_, err := p.Get()
 		if !errors.Is(err, errTest) {
@@ -298,7 +298,7 @@ func TestPool(t *testing.T) {
 	})
 
 	t.Run("use default algo", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+		alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 		p := pool.Make(alloc, nil, pool.Options{})
 
 		var bs [][]byte
@@ -323,7 +323,7 @@ func TestPool(t *testing.T) {
 	})
 
 	t.Run("drop", func(t *testing.T) {
-		alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+		alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 		p := pool.Make(alloc, nil, pool.Options{})
 		_, err := p.Get()
 		if err != nil {
@@ -339,5 +339,389 @@ func TestPool(t *testing.T) {
 		p.Drop()
 		s = p.Stats()
 		e = pool.Stats{Alloc: 1, Get: 1, Drop: 1}
+		if s != e {
+			t.Fatal(s)
+		}
+	})
+
+	t.Run("key", func(t *testing.T) {
+		t.Run("distinct", func(t *testing.T) {
+			alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
+			p := pool.Make(alloc, nil, pool.Options{Algo: pool.NoShrink()})
+			f, err := p.GetKey("foo")
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if len(f) != 1<<9 {
+				t.Fatal(len(f))
+			}
+
+			b, err := p.GetKey("bar")
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if len(b) != 1<<9 {
+				t.Fatal(len(b))
+			}
+
+			p.PutKey("foo", f)
+			f, err = p.GetKey("foo")
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if len(f) != 1<<9 {
+				t.Fatal(len(f))
+			}
+
+			b, err = p.GetKey("bar")
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if len(b) != 1<<9 {
+				t.Fatal(len(b))
+			}
+
+			s := p.Stats()
+			e := pool.Stats{Alloc: 3, Get: 4, Active: 3, Put: 1}
+			if s != e {
+				t.Fatal(s)
+			}
+		})
+
+		t.Run("get", func(t *testing.T) {
+			t.Run("empty", func(t *testing.T) {
+				alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
+				p := pool.Make(alloc, nil, pool.Options{Algo: pool.NoShrink()})
+				b, err := p.GetKey("foo")
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				if len(b) != 1<<9 {
+					t.Fatal(len(b))
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Alloc: 1, Get: 1, Active: 1}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("pooled own", func(t *testing.T) {
+				alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
+				p := pool.Make(alloc, nil, pool.Options{Algo: pool.NoShrink()})
+				b, err := p.GetKey("foo")
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				if len(b) != 1<<9 {
+					t.Fatal(len(b))
+				}
+
+				p.PutKey("foo", b)
+				b, err = p.GetKey("foo")
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				if len(b) != 1<<9 {
+					t.Fatal(len(b))
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Alloc: 1, Get: 2, Put: 1, Active: 1}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("pooled foreign", func(t *testing.T) {
+				alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
+				p := pool.Make(alloc, nil, pool.Options{Algo: pool.NoShrink()})
+				p.PutKey("foo", make([]byte, 1<<6))
+				b, err := p.GetKey("foo")
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				if len(b) != 1<<6 {
+					t.Fatal(len(b))
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Get: 1, Put: 1, Active: 1}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("own alloc not available", func(t *testing.T) {
+				p := pool.Make[[]byte](nil, nil, pool.Options{Algo: pool.NoShrink()})
+				_, err := p.GetKey("foo")
+				if !errors.Is(err, pool.ErrNoItems) {
+					t.Fatal(err)
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Get: 1}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("foreign no alloc", func(t *testing.T) {
+				p := pool.Make[[]byte](nil, nil, pool.Options{Algo: pool.NoShrink()})
+				p.PutKey("foo", make([]byte, 1<<6))
+				b, err := p.GetKey("foo")
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				if len(b) != 1<<6 {
+					t.Fatal(len(b))
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Get: 1, Put: 1, Active: 1}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("multiple", func(t *testing.T) {
+				alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
+				p := pool.Make(alloc, nil, pool.Options{Algo: pool.NoShrink()})
+				b0, err := p.GetKey("foo")
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				b1, err := p.GetKey("foo")
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				b2, err := p.GetKey("foo")
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				for _, b := range [][]byte{b0, b1, b2} {
+					if len(b) != 1<<9 {
+						t.Fatal(len(b))
+					}
+				}
+
+				p.PutKey("foo", b0)
+				p.PutKey("foo", b1)
+				p.PutKey("foo", b2)
+				b, err := p.GetKeyN("foo", 3)
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				if len(b) != 3 {
+					t.Fatal(len(b))
+				}
+
+				for i := range b {
+					if len(b[i]) != 1<<9 {
+						t.Fatal(len(b))
+					}
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Alloc: 3, Get: 4, Put: 3, Active: 3}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+		})
+
+		t.Run("put and release", func(t *testing.T) {
+			t.Run("foreign not empty", func(t *testing.T) {
+				p := pool.Make[[]byte](nil, nil, pool.Options{Algo: pool.NoShrink()})
+				p.PutKey("foo", make([]byte, 1<<6))
+				p.PutKey("foo", make([]byte, 1<<6))
+				b, err := p.GetKey("foo")
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				if len(b) != 1<<6 {
+					t.Fatal(len(b))
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Get: 1, Put: 2, Active: 1, Idle: 1}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("release on put no free", func(t *testing.T) {
+				p := pool.Make[[]byte](nil, nil, pool.Options{Algo: pool.Max(2)})
+				p.PutKey("foo", make([]byte, 1<<9))
+				p.PutKey("foo", make([]byte, 1<<9))
+				p.PutKey("foo", make([]byte, 1<<9))
+				s := p.Stats()
+				e := pool.Stats{Put: 3, Idle: 2, Free: 1}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("release on put with free", func(t *testing.T) {
+				var freeCount int
+				f := func([]byte) { freeCount++ }
+				p := pool.Make(nil, f, pool.Options{Algo: pool.Max(2)})
+				p.PutKey("foo", make([]byte, 1<<9))
+				p.PutKey("foo", make([]byte, 1<<9))
+				p.PutKey("foo", make([]byte, 1<<9))
+				if freeCount != 1 {
+					t.Fatal(freeCount)
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Put: 3, Idle: 2, Free: 1}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("release all no free", func(t *testing.T) {
+				p := pool.Make[[]byte](nil, nil, pool.Options{Algo: pool.NoShrink()})
+				p.PutKey("foo", make([]byte, 1<<9))
+				p.PutKey("foo", make([]byte, 1<<9))
+				p.PutKey("foo", make([]byte, 1<<9))
+				p.Free()
+				s := p.Stats()
+				e := pool.Stats{Put: 3, Idle: 0, Free: 3}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("release all with free", func(t *testing.T) {
+				var freeCount int
+				f := func([]byte) { freeCount++ }
+				p := pool.Make[[]byte](nil, f, pool.Options{Algo: pool.NoShrink()})
+				p.PutKey("foo", make([]byte, 1<<9))
+				p.PutKey("foo", make([]byte, 1<<9))
+				p.PutKey("foo", make([]byte, 1<<9))
+				p.Free()
+				if freeCount != 3 {
+					t.Fatal(freeCount)
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Put: 3, Idle: 0, Free: 3}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("release on timeout no free", func(t *testing.T) {
+				c := times.Test()
+				b := syncbus.New(time.Second)
+
+				var o pool.Options
+				o.Algo = pool.Timeout(3 * time.Millisecond)
+				o.Testing.Clock = c
+				o.Testing.Bus = b
+				p := pool.Make[[]byte](nil, nil, o)
+				p.PutKey("foo", make([]byte, 1<<9))
+				if err := b.Wait("background-job-waiting"); err != nil {
+					t.Fatal(err)
+				}
+
+				c.Pass(2 * time.Millisecond)
+				p.PutKey("foo", make([]byte, 1<<9))
+				c.Pass(2 * time.Millisecond)
+				if err := b.Wait("free-idle-done"); err != nil {
+					t.Fatal(err)
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Put: 2, Idle: 1, Free: 1}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+
+			t.Run("release on timeout with free", func(t *testing.T) {
+				var freeCount int
+				f := func([]byte) { freeCount++ }
+				c := times.Test()
+				b := syncbus.New(time.Second)
+
+				var o pool.Options
+				o.Algo = pool.Timeout(3 * time.Millisecond)
+				o.Testing.Clock = c
+				o.Testing.Bus = b
+				p := pool.Make[[]byte](nil, f, o)
+				p.PutKey("foo", make([]byte, 1<<9))
+				if err := b.Wait("background-job-waiting"); err != nil {
+					t.Fatal(err)
+				}
+
+				c.Pass(2 * time.Millisecond)
+				p.PutKey("foo", make([]byte, 1<<9))
+				c.Pass(2 * time.Millisecond)
+				if err := b.Wait("free-idle-done"); err != nil {
+					t.Fatal(err)
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Put: 2, Idle: 1, Free: 1}
+				if s != e || freeCount != 1 {
+					t.Fatal(s, freeCount)
+				}
+			})
+
+			t.Run("multiple", func(t *testing.T) {
+				p := pool.Make[[]byte](nil, nil, pool.Options{Algo: pool.NoShrink()})
+				p.PutKey(
+					"foo",
+					make([]byte, 1<<6),
+					make([]byte, 1<<6),
+					make([]byte, 1<<6),
+				)
+				b, err := p.GetKey("foo")
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				if len(b) != 1<<6 {
+					t.Fatal(len(b))
+				}
+
+				s := p.Stats()
+				e := pool.Stats{Get: 1, Put: 1, Active: 1, Idle: 2}
+				if s != e {
+					t.Fatal(s)
+				}
+			})
+		})
+
+		t.Run("load", func(t *testing.T) {
+			p := pool.Make[[]byte](nil, nil, pool.Options{Algo: pool.NoShrink()})
+			l := make([][]byte, 9)
+			for i := 0; i < len(l); i++ {
+				l[i] = make([]byte, 1<<9)
+			}
+
+			p.LoadKey("foo", l)
+			s := p.Stats()
+			if s.Idle != 9 {
+				t.Fatal(s)
+			}
+		})
 	})
 }

scenario_test.go

@@ -96,7 +96,7 @@ func (s stats) push(a pool.Stats) {
 func testScenario(t *testing.T, o scenarioOptions, step scenarioStep, verify verifyScenario) {
 	for _, n := range o.exclude {
 		if strings.HasSuffix(t.Name(), n) {
-			t.Skip()
+			t.Skip("test excluded")
 		}
 	}
 
@@ -119,7 +119,7 @@ func testScenario(t *testing.T, o scenarioOptions, step scenarioStep, verify ver
 		clock = testClock
 	}
 
-	alloc := func() ([]byte, error) { return make([]byte, 1<<9), nil }
+	alloc := func(string) ([]byte, error) { return make([]byte, 1<<9), nil }
 
 	var po pool.Options
 	po.Algo = o.algo