diff --git a/cmd/cgat/main.go b/cmd/cgat/main.go
index 20108d1978150527696a92dbf065d2c383eaea22..6058bc691b9cbb5ad01422795d8b90d3255ba39e 100644
--- a/cmd/cgat/main.go
+++ b/cmd/cgat/main.go
@@ -8,7 +8,7 @@ import (
"pggat2/lib/bouncer/backends/v0"
"pggat2/lib/bouncer/bouncers/v0"
"pggat2/lib/bouncer/frontends/v0"
- "pggat2/lib/middleware/middlewares/unread"
+ "pggat2/lib/middleware/middlewares/onebuffer"
"pggat2/lib/middleware/middlewares/unterminate"
"pggat2/lib/rob"
"pggat2/lib/rob/schedulers/v2"
@@ -61,15 +61,16 @@ func main() {
client := zio.MakeReadWriter(conn)
ut := unterminate.MakeUnterminate(&client)
frontends.Accept(ut)
+ ob := onebuffer.MakeOnebuffer(&client)
done := make(chan struct{})
defer close(done)
for {
- u, err := unread.NewUnread(&ut)
+ err = ob.Buffer()
if err != nil {
break
}
source.Schedule(job{
- client: u,
+ client: &ob,
done: done,
}, 0)
<-done
diff --git a/lib/README.md b/lib/README.md
index 3a0c0beb933338280d771c815c771bbd7942d8be..dd8a64fc30b6af2bcb4e336fad9a48cfdf54a953 100644
--- a/lib/README.md
+++ b/lib/README.md
@@ -6,7 +6,7 @@ In general, the top level folder will hold the interface. A subfolder (generally
## auth
All authentication functions. Protocol unspecific.
-## bounce
+## bouncer
All routing: accepting frontends, backends, and handling transactions.
## middleware
@@ -15,14 +15,14 @@ Intercept packets and perform operations on them
## perror
Special postgres error types
-## pnet
-Zero allocation network handling
-
-### pnet/packet
-Packet reading/writing helpers
-
## rob
A fair-share scheduler
## util
Project generic helper structures and functions
+
+## zap
+Zero allocation packet handling
+
+### zap/packets
+Packet reading/writing helpers
diff --git a/lib/middleware/middlewares/onebuffer/onebuffer.go b/lib/middleware/middlewares/onebuffer/onebuffer.go
new file mode 100644
index 0000000000000000000000000000000000000000..3ecca273fe48d68b44f96ee178c4a7d71d04efa0
--- /dev/null
+++ b/lib/middleware/middlewares/onebuffer/onebuffer.go
@@ -0,0 +1,57 @@
+package onebuffer
+
+import (
+ "pggat2/lib/util/decorator"
+ "pggat2/lib/zap"
+)
+
+type Onebuffer struct {
+ noCopy decorator.NoCopy
+
+ in zap.In
+ read bool
+ zap.ReadWriter
+}
+
+func MakeOnebuffer(inner zap.ReadWriter) Onebuffer {
+ return Onebuffer{
+ read: true,
+ ReadWriter: inner,
+ }
+}
+
+func (T *Onebuffer) Buffer() error {
+ if !T.read {
+ panic("a packet is already buffered in the Onebuffer!")
+ }
+ var err error
+ T.in, err = T.ReadWriter.Read()
+ T.read = false
+ return err
+}
+
+func (T *Onebuffer) BufferUntyped() error {
+ if !T.read {
+ panic("a packet is already buffered in the Onebuffer!")
+ }
+ var err error
+ T.in, err = T.ReadWriter.ReadUntyped()
+ T.read = false
+ return err
+}
+
+func (T *Onebuffer) Read() (zap.In, error) {
+ if !T.read {
+ T.read = true
+ return T.in, nil
+ }
+ return T.ReadWriter.Read()
+}
+
+func (T *Onebuffer) ReadUntyped() (zap.In, error) {
+ if !T.read {
+ T.read = true
+ return T.in, nil
+ }
+ return T.ReadWriter.ReadUntyped()
+}
diff --git a/lib/middleware/middlewares/unread/unread.go b/lib/middleware/middlewares/unread/unread.go
deleted file mode 100644
index c5a860215c675e2607fe530c25c06c8d1316db7f..0000000000000000000000000000000000000000
--- a/lib/middleware/middlewares/unread/unread.go
+++ /dev/null
@@ -1,51 +0,0 @@
-package unread
-
-import (
- "pggat2/lib/zap"
-)
-
-type Unread struct {
- in zap.In
- read bool
- zap.ReadWriter
-}
-
-func NewUnread(inner zap.ReadWriter) (*Unread, error) {
- in, err := inner.Read()
- if err != nil {
- return nil, err
- }
- return &Unread{
- in: in,
- ReadWriter: inner,
- }, nil
-}
-
-func NewUnreadUntyped(inner zap.ReadWriter) (*Unread, error) {
- in, err := inner.ReadUntyped()
- if err != nil {
- return nil, err
- }
- return &Unread{
- in: in,
- ReadWriter: inner,
- }, nil
-}
-
-func (T *Unread) Read() (zap.In, error) {
- if !T.read {
- T.read = true
- return T.in, nil
- }
- return T.ReadWriter.Read()
-}
-
-func (T *Unread) ReadUntyped() (zap.In, error) {
- if !T.read {
- T.read = true
- return T.in, nil
- }
- return T.ReadWriter.ReadUntyped()
-}
-
-var _ zap.ReadWriter = (*Unread)(nil)
diff --git a/lib/util/rbtree/rbtree.go b/lib/util/rbtree/rbtree.go
index 5c93f00d08189a3337f6ea9de99e2b1fa9ae3238..f22265cb43f656584a54e60cfc28bbb96b2138c1 100644
--- a/lib/util/rbtree/rbtree.go
+++ b/lib/util/rbtree/rbtree.go
@@ -9,6 +9,22 @@ type order interface {
// RBTree is a left-leaning red-black BST
type RBTree[K order, V any] struct {
root *node[K, V]
+ // pool of node to reuse memory
+ pool []*node[K, V]
+}
+
+func (T *RBTree[K, V]) free(n *node[K, V]) {
+ T.pool = append(T.pool, n)
+}
+
+func (T *RBTree[K, V]) alloc() *node[K, V] {
+ if len(T.pool) > 0 {
+ v := T.pool[len(T.pool)-1]
+ T.pool = T.pool[:len(T.pool)-1]
+ *v = node[K, V]{}
+ return v
+ }
+ return new(node[K, V])
}
func (T *RBTree[K, V]) Get(key K) (V, bool) {
@@ -32,11 +48,11 @@ func (T *RBTree[K, V]) Set(key K, value V) {
func (T *RBTree[K, V]) put(n *node[K, V], key K, value V) *node[K, V] {
if n == nil {
- return &node[K, V]{
- key: key,
- value: value,
- color: red,
- }
+ n = T.alloc()
+ n.key = key
+ n.value = value
+ n.color = red
+ return n
}
if key > n.key {
@@ -76,6 +92,7 @@ func (T *RBTree[K, V]) delete(n *node[K, V], key K) *node[K, V] {
n = T.rotateRight(n)
}
if key == n.key && n.right == nil {
+ T.free(n)
return nil
}
if n.right.getColor() == black && n.right.left.getColor() == black {
@@ -95,6 +112,7 @@ func (T *RBTree[K, V]) delete(n *node[K, V], key K) *node[K, V] {
func (T *RBTree[K, V]) deleteMin(n *node[K, V]) *node[K, V] {
if n.left == nil {
+ T.free(n)
return nil
}
diff --git a/lib/util/rbtree/rbtree_test.go b/lib/util/rbtree/rbtree_test.go
index 7227adaf7cac0eaa036f029610b11f43010cb5a4..4878b46dea81c05b8877d2b34b479f6d4b06d4b9 100644
--- a/lib/util/rbtree/rbtree_test.go
+++ b/lib/util/rbtree/rbtree_test.go
@@ -101,11 +101,37 @@ func TestRBTree_Stress(t *testing.T) {
tree := new(RBTree[int, int])
+ existing := make(map[int]struct{})
+
for i := 0; i < n; i++ {
- k := rand.Int()
- v := rand.Int()
- tree.Set(k, v)
- tree.Set(rand.Int(), rand.Int())
- tree.Delete(k)
+ k1 := rand.Int()
+ v1 := rand.Int()
+ k2 := rand.Int()
+ v2 := rand.Int()
+ tree.Set(k1, v1)
+ tree.Set(k2, v2)
+ tree.Delete(k1)
+ delete(existing, k1)
+ existing[k2] = struct{}{}
+ }
+
+ // make sure all values we expect are there
+ for k := range existing {
+ if _, ok := tree.Get(k); !ok {
+ t.Error("expected value to exist")
+ }
+ }
+
+ // iterate through to make sure there aren't any extra values
+ prev := 0
+ for k, v, ok := tree.Min(); ok; k, v, ok = tree.Next(k) {
+ if k < prev {
+ t.Error("expected tree to be correctly ordered")
+ }
+ prev = k
+ if _, ok = existing[k]; !ok {
+ t.Error("unexpected value in tree")
+ }
+ _ = v
}
}