1 files changed, 53 insertions, 0 deletions
diff --git a/llm/pool.go b/llm/pool.go
new file mode 100644
index 0000000..294b64f
--- /dev/null
+++ b/llm/pool.go
@@ -0,0 +1,53 @@
+package llm
+
+import "sync"
+
+type pool[K comparable] struct {
+	devices map[K]bool
+	mutex   sync.Mutex
+	cond    *sync.Cond
+}
+
+func newPool[K comparable](devices ...K) *pool[K] {
+	p := &pool[K]{
+		devices: make(map[K]bool),
+	}
+
+	for _, d := range devices {
+		p.devices[d] = true
+	}
+
+	p.cond = sync.NewCond(&p.mutex)
+
+	return p
+}
+
+func (p *pool[K]) Len() int {
+	return len(p.devices)
+}
+
+func (p *pool[K]) Acquire() K {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+
+	for {
+		for k, v := range p.devices {
+			if v {
+				p.devices[k] = false
+
+				return k
+			}
+		}
+
+		p.cond.Wait()
+	}
+}
+
+func (p *pool[K]) Release(device K) {
+	p.mutex.Lock()
+
+	p.devices[device] = true
+
+	p.cond.Signal()
+	p.mutex.Unlock()
+}