From 8a90457823fdf9735f09bc05476a83aaa377ecc0 Mon Sep 17 00:00:00 2001
From: Varun Patil <radialapps@gmail.com>
Date: Tue, 24 Dec 2024 23:28:24 -0800
Subject: [PATCH] pit: add entry and node pools

---
 fw/table/pit-cs-tree.go | 42 +++++++++++++++++++++++++++++++----------
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/fw/table/pit-cs-tree.go b/fw/table/pit-cs-tree.go
index 0270498b..6788f8a0 100644
--- a/fw/table/pit-cs-tree.go
+++ b/fw/table/pit-cs-tree.go
@@ -2,6 +2,7 @@ package table
 
 import (
 	"math/rand"
+	"sync"
 	"time"
 
 	"github.com/named-data/ndnd/fw/core"
@@ -14,6 +15,24 @@ const expiredPitTickerInterval = 100 * time.Millisecond
 
 type OnPitExpiration func(PitEntry)
 
+var nameTreePitEntryPool = sync.Pool{
+	New: func() interface{} {
+		entry := new(nameTreePitEntry)
+		entry.inRecords = make(map[uint64]*PitInRecord)
+		entry.outRecords = make(map[uint64]*PitOutRecord)
+		return entry
+	},
+}
+
+var pitCsTreeNodePool = sync.Pool{
+	New: func() interface{} {
+		node := new(pitCsTreeNode)
+		node.children = make(map[uint64]*pitCsTreeNode)
+		node.pitEntries = make([]*nameTreePitEntry, 0, 1)
+		return node
+	},
+}
+
 // PitCsTree represents a PIT-CS implementation that uses a name tree
 type PitCsTree struct {
 	basePitCsTable
@@ -147,19 +166,19 @@ func (p *PitCsTree) InsertInterest(interest ndn.Interest, hint enc.Name, inFace
 
 	if entry == nil {
 		p.nPitEntries++
-		entry = new(nameTreePitEntry)
+		entry = nameTreePitEntryPool.Get().(*nameTreePitEntry)
 		entry.node = node
 		entry.pitCsTable = p
 		entry.encname = name
 		entry.canBePrefix = interest.CanBePrefix()
 		entry.mustBeFresh = interest.MustBeFresh()
 		entry.forwardingHintNew = hint
-		entry.inRecords = make(map[uint64]*PitInRecord)
-		entry.outRecords = make(map[uint64]*PitOutRecord)
+		clear(entry.inRecords)
+		clear(entry.outRecords)
 		entry.satisfied = false
-		node.pitEntries = append(node.pitEntries, entry)
 		entry.token = p.generateNewPitToken()
 		entry.pqItem = nil
+		node.pitEntries = append(node.pitEntries, entry)
 		p.pitTokenMap[entry.token] = entry
 	}
 
@@ -192,6 +211,7 @@ func (p *PitCsTree) RemoveInterest(pitEntry PitEntry) bool {
 			}
 			p.nPitEntries--
 			delete(p.pitTokenMap, e.token)
+			nameTreePitEntryPool.Put(e)
 			return true
 		}
 	}
@@ -337,14 +357,15 @@ func (p *pitCsTreeNode) findLongestPrefixEntryEnc(name enc.Name) *pitCsTreeNode
 func (p *pitCsTreeNode) fillTreeToPrefixEnc(name enc.Name) *pitCsTreeNode {
 	curNode := p.findLongestPrefixEntryEnc(name)
 	for depth := curNode.depth + 1; depth <= len(name); depth++ {
-		newNode := new(pitCsTreeNode)
-		var temp = At(name, depth-1)
-		newNode.component = &temp
+		comp := At(name, depth-1)
+
+		newNode := pitCsTreeNodePool.Get().(*pitCsTreeNode)
+		newNode.component = &comp
 		newNode.depth = depth
 		newNode.parent = curNode
-		newNode.children = make(map[uint64]*pitCsTreeNode)
-
-		curNode.children[newNode.component.Hash()] = newNode
+		clear(newNode.children)
+		clear(newNode.pitEntries)
+		curNode.children[comp.Hash()] = newNode
 		curNode = newNode
 	}
 	return curNode
@@ -358,6 +379,7 @@ func (p *pitCsTreeNode) pruneIfEmpty() {
 	for curNode := p; curNode.parent != nil && curNode.getChildrenCount() == 0 &&
 		len(curNode.pitEntries) == 0 && curNode.csEntry == nil; curNode = curNode.parent {
 		delete(curNode.parent.children, curNode.component.Hash())
+		pitCsTreeNodePool.Put(curNode)
 	}
 }