Refactor pool stats

cmd/restic/main.go

@@ -178,5 +178,6 @@ func main() {
 		os.Exit(1)
 	}
 
-	// restic.PoolAlloc()
+	// this prints some statistics for memory management using the debug package
+	restic.PoolAlloc()
 }

pools.go

@@ -1,103 +1,113 @@
 package restic
 
-import "sync"
+import (
+	"sync"
+
+	"github.com/restic/restic/debug"
+)
+
+type poolStats struct {
+	m    sync.Mutex
+	mget map[string]int
+	mput map[string]int
+	mmax map[string]int
+
+	get int
+	put int
+	max int
+}
+
+func (s *poolStats) Get(k string) {
+	s.m.Lock()
+	defer s.m.Unlock()
+
+	s.get += 1
+	cur := s.get - s.put
+	if cur > s.max {
+		s.max = cur
+	}
+
+	if k != "" {
+		if _, ok := s.mget[k]; !ok {
+			s.mget[k] = 0
+			s.mput[k] = 0
+			s.mmax[k] = 0
+		}
+
+		s.mget[k]++
+
+		cur = s.mget[k] - s.mput[k]
+		if cur > s.mmax[k] {
+			s.mmax[k] = cur
+		}
+	}
+}
+
+func (s *poolStats) Put(k string) {
+	s.m.Lock()
+	defer s.m.Unlock()
+
+	s.put += 1
+
+	if k != "" {
+		s.mput[k]++
+	}
+}
+
+func newPoolStats() *poolStats {
+	return &poolStats{
+		mget: make(map[string]int),
+		mput: make(map[string]int),
+		mmax: make(map[string]int),
+	}
+}
 
 var (
 	chunkPool = sync.Pool{New: newChunkBuf}
 	nodePool  = sync.Pool{New: newNode}
+
+	chunkStats = newPoolStats()
+	nodeStats  = newPoolStats()
 )
 
-type alloc_stats struct {
-	m         sync.Mutex
-	alloc_map map[string]int
-	free_map  map[string]int
-	alloc     int
-	free      int
-	new       int
-	all       int
-	max       int
-}
-
-var (
-	chunk_stats alloc_stats
-	node_stats  alloc_stats
-)
-
-func init() {
-	chunk_stats.alloc_map = make(map[string]int)
-	chunk_stats.free_map = make(map[string]int)
-}
-
 func newChunkBuf() interface{} {
-	chunk_stats.m.Lock()
-	chunk_stats.new += 1
-	chunk_stats.m.Unlock()
-
 	// create buffer for iv, data and hmac
 	return make([]byte, maxCiphertextSize)
 }
 
 func newNode() interface{} {
-	node_stats.m.Lock()
-	node_stats.new += 1
-	node_stats.m.Unlock()
-
 	// create buffer for iv, data and hmac
 	return new(Node)
 }
 
 func GetChunkBuf(s string) []byte {
-	chunk_stats.m.Lock()
-	if _, ok := chunk_stats.alloc_map[s]; !ok {
-		chunk_stats.alloc_map[s] = 0
-	}
-	chunk_stats.alloc_map[s] += 1
-	chunk_stats.all += 1
-	if chunk_stats.all > chunk_stats.max {
-		chunk_stats.max = chunk_stats.all
-	}
-	chunk_stats.m.Unlock()
-
+	chunkStats.Get(s)
 	return chunkPool.Get().([]byte)
 }
 
 func FreeChunkBuf(s string, buf []byte) {
-	chunk_stats.m.Lock()
-	if _, ok := chunk_stats.free_map[s]; !ok {
-		chunk_stats.free_map[s] = 0
-	}
-	chunk_stats.free_map[s] += 1
-	chunk_stats.all -= 1
-	chunk_stats.m.Unlock()
-
+	chunkStats.Put(s)
 	chunkPool.Put(buf)
 }
 
 func GetNode() *Node {
-	node_stats.m.Lock()
-	node_stats.alloc += 1
-	node_stats.all += 1
-	if node_stats.all > node_stats.max {
-		node_stats.max = node_stats.all
-	}
-	node_stats.m.Unlock()
+	nodeStats.Get("")
 	return nodePool.Get().(*Node)
 }
 
 func FreeNode(n *Node) {
-	node_stats.m.Lock()
-	node_stats.all -= 1
-	node_stats.free += 1
-	node_stats.m.Unlock()
+	nodeStats.Put("")
 	nodePool.Put(n)
 }
 
 func PoolAlloc() {
-	// fmt.Fprintf(os.Stderr, "alloc max: %d, new: %d\n", chunk_stats.max, chunk_stats.new)
-	// for k, v := range chunk_stats.alloc_map {
-	// 	fmt.Fprintf(os.Stderr, "alloc[%s] %d, free %d diff: %d\n", k, v, chunk_stats.free_map[k], v-chunk_stats.free_map[k])
-	// }
+	debug.Log("pools.PoolAlloc", "pool stats for chunk: get %d, put %d, diff %d, max %d\n", chunkStats.get, chunkStats.put, chunkStats.get-chunkStats.put, chunkStats.max)
+	for k, v := range chunkStats.mget {
+		debug.Log("pools.PoolAlloc", "pool stats for chunk[%s]: get %d, put %d, diff %d, max %d\n", k, v, chunkStats.mput[k], v-chunkStats.mput[k], chunkStats.mmax[k])
+	}
 
-	// fmt.Fprintf(os.Stderr, "nodes alloc max: %d, new: %d\n", node_stats.max, node_stats.new)
-	// fmt.Fprintf(os.Stderr, "alloc %d, free %d diff: %d\n", node_stats.alloc, node_stats.free, node_stats.alloc-node_stats.free)
+	debug.Log("pools.PoolAlloc", "pool stats for node: get %d, put %d, diff %d, max %d\n", nodeStats.get, nodeStats.put, nodeStats.get-nodeStats.put, nodeStats.max)
+	for k, v := range nodeStats.mget {
+		debug.Log("pools.PoolAlloc", "pool stats for node[%s]: get %d, put %d, diff %d, max %d\n", k, v, nodeStats.mput[k], v-nodeStats.mput[k], nodeStats.mmax[k])
+	}
 }