Add LRU limit for channels

Author: floatdrop

dynamic.go

@@ -1,36 +1,117 @@
 package demux
 
-import "sync"
+import (
+	"container/list"
+	"math"
+	"sync"
+)
+
+type channelInfo[T any, K comparable] struct {
+	ch      chan T
+	key     K
+	element *list.Element
+}
+
+// Config holds configuration for the dynamic demultiplexer
+type Config struct {
+	maxChannels int
+	bufferSize  int
+}
+
+// Option is a functional option for configuring the demultiplexer
+type Option func(*Config)
+
+// WithMaxChannels sets the maximum number of concurrent channels
+func WithMaxChannels(max int) Option {
+	return func(c *Config) {
+		if max > 0 {
+			c.maxChannels = max
+		}
+	}
+}
+
+// WithBufferSize sets the buffer size for each channel
+func WithBufferSize(size int) Option {
+	return func(c *Config) {
+		if size >= 0 {
+			c.bufferSize = size
+		}
+	}
+}
+
+// defaultConfig returns the default configuration
+func defaultConfig() Config {
+	return Config{
+		maxChannels: math.MaxInt, // uncapped by default
+		bufferSize:  0,           // unbuffered by default
+	}
+}
 
 // Dynamic creates dynamic demultiplexer that routes items from 'in' based on keys returned by 'keyFunc'.
 // For each unique key, a new goroutine is spawned running 'consumeFunc'.
 // Each consumeFunc receives a channel that delivers values matching its key.
+// When maxChannels limit is reached, least recently used channels are evicted.
 func Dynamic[T any, K comparable](
 	in <-chan T,
 	keyFunc func(T) K,
 	consumeFunc func(K, <-chan T),
+	opts ...Option,
 ) {
-	outChans := make(map[K]chan T)
+	// Apply options
+	config := defaultConfig()
+	for _, opt := range opts {
+		opt(&config)
+	}
+
+	outChans := make(map[K]*channelInfo[T, K])
+	lru := list.New()
 	var wg sync.WaitGroup
 
 	for t := range in {
 		key := keyFunc(t)
-		ch, exists := outChans[key]
-		if !exists {
-			ch = make(chan T)
-			outChans[key] = ch
+
+		info, exists := outChans[key]
+
+		if exists {
+			// Move to front (most recently used)
+			lru.MoveToFront(info.element)
+		} else {
+			// Check if we need to evict
+			if len(outChans) >= config.maxChannels {
+				// Evict least recently used
+				oldest := lru.Back()
+				if oldest != nil {
+					oldInfo := oldest.Value.(*channelInfo[T, K])
+					close(oldInfo.ch)
+					delete(outChans, oldInfo.key)
+					lru.Remove(oldest)
+				}
+			}
+
+			// Create new channel with configured buffer size
+			ch := make(chan T, config.bufferSize)
+			element := lru.PushFront(nil)
+			info = &channelInfo[T, K]{
+				ch:      ch,
+				key:     key,
+				element: element,
+			}
+			element.Value = info
+			outChans[key] = info
 
 			wg.Add(1)
-			go func() {
+			go func(k K, c <-chan T) {
 				defer wg.Done()
-				consumeFunc(key, ch)
-			}()
+				consumeFunc(k, c)
+			}(key, ch)
 		}
-		ch <- t
+
+		info.ch <- t
 	}
 
-	for _, ch := range outChans {
-		close(ch)
+	// Close all remaining channels
+	for _, info := range outChans {
+		close(info.ch)
 	}
 
 	wg.Wait()