From e7b2d0e7769717035e214626ecd286c4b515cabe Mon Sep 17 00:00:00 2001
From: swdee <shanon@bokumail.com>
Date: Mon, 7 Jul 2025 15:29:52 +1200
Subject: [PATCH] added BatchPool functions to have a pool of batches

---
 batchpool.go | 75 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 75 insertions(+)
 create mode 100644 batchpool.go

diff --git a/batchpool.go b/batchpool.go
new file mode 100644
index 0000000..a09fbe7
--- /dev/null
+++ b/batchpool.go
@@ -0,0 +1,75 @@
+package rknnlite
+
+import (
+	"sync"
+)
+
+// BatchPool is a pool of batches
+type BatchPool struct {
+	// pool of batches
+	batches chan *Batch
+	// size of pool
+	size  int
+	close sync.Once
+}
+
+// NewBatchPool returns a pool of Batches
+func NewBatchPool(size int, rt *Runtime) *BatchPool {
+
+	p := &BatchPool{
+		batches: make(chan *Batch, size),
+		size:    size,
+	}
+
+	batchSize := int(rt.InputAttrs()[0].Dims[0])
+	width := int(rt.InputAttrs()[0].Dims[1])
+	height := int(rt.InputAttrs()[0].Dims[2])
+	channels := int(rt.InputAttrs()[0].Dims[3])
+	inputType := rt.GetInputTypeFloat32()
+
+	// create batch pool to be the same size as the runtime pool
+	for i := 0; i < size; i++ {
+		batch := NewBatch(
+			batchSize,
+			height,
+			width,
+			channels,
+			inputType,
+		)
+
+		// attach to pool
+		p.Return(batch)
+	}
+
+	return p
+}
+
+// Gets a batch from the pool
+func (p *BatchPool) Get() *Batch {
+	return <-p.batches
+}
+
+// Return a batch to the pool
+func (p *BatchPool) Return(batch *Batch) {
+
+	batch.Clear()
+
+	select {
+	case p.batches <- batch:
+	default:
+		// pool is full or closed
+	}
+}
+
+// Close the pool and all batches in it
+func (p *BatchPool) Close() {
+	p.close.Do(func() {
+		// close channel
+		close(p.batches)
+
+		// close all runtimes
+		for next := range p.batches {
+			_ = next.Close()
+		}
+	})
+}