Heap pool

src/core/cycle_detector.go

@@ -16,7 +16,7 @@ func (c *cycleDetector) Check() *errCycle {
 	if c.stopped {
 		return nil
 	}
-	log.Debug("Running cycle detection...")
+	//log.Debug("Running cycle detection...")
 	complete := map[*BuildTarget]struct{}{}
 	partial := map[*BuildTarget]struct{}{}
 
@@ -55,12 +55,12 @@ func (c *cycleDetector) Check() *errCycle {
 		}
 		if _, present := complete[target]; !present {
 			if cycle, _ := visit(target); cycle != nil {
-				log.Debug("Cycle detection complete, cycle found: %s", cycle)
+				//log.Debug("Cycle detection complete, cycle found: %s", cycle)
 				return &errCycle{Cycle: cycle}
 			}
 		}
 	}
-	log.Debug("Cycle detection complete, no cycles found")
+	//log.Debug("Cycle detection complete, no cycles found")
 	return nil
 }
 

src/parse/asp/benchmark_test.go

@@ -6,10 +6,14 @@ import (
 	"fmt"
 	"github.com/thought-machine/please/rules"
 	"github.com/thought-machine/please/src/core"
+	"github.com/thought-machine/please/src/parse/asp/heap"
 	"os"
 	"path/filepath"
+	"runtime"
+	"runtime/debug"
 	"sync"
 	"testing"
+	"time"
 )
 
 var code = `
@@ -75,7 +79,6 @@ go_test(
 `
 
 func BenchmarkParse(b *testing.B) {
-	arena.NewArena().Free()
 	b.ReportAllocs()
 
 	for i := 0; i < b.N; i++ {
@@ -84,7 +87,6 @@ func BenchmarkParse(b *testing.B) {
 }
 
 func BenchmarkParseWithArena(b *testing.B) {
-	arena.NewArena().Free()
 	b.ReportAllocs()
 
 	for i := 0; i < b.N; i++ {
@@ -95,6 +97,7 @@ func BenchmarkParseWithArena(b *testing.B) {
 func parseInParallel(threads, repeats int, useArena bool) {
 	wg := new(sync.WaitGroup)
 	wg.Add(threads)
+	pool := heap.NewPool(threads, -1, time.Second)
 	for j := 0; j < threads; j++ {
 		go func() {
 			for k := 0; k < repeats; k++ {
@@ -103,9 +106,9 @@ func parseInParallel(threads, repeats int, useArena bool) {
 					parseFileInput(r, nil)
 					continue
 				}
-				a := arena.NewArena()
-				parseFileInput(r, a)
-				a.Free()
+				heap := pool.Get()
+				parseFileInput(r, heap.Arena)
+				pool.Return(heap)
 			}
 			wg.Done()
 		}()
@@ -116,19 +119,17 @@ func parseInParallel(threads, repeats int, useArena bool) {
 func BenchmarkParseAndInterpretWithArena(b *testing.B) {
 	b.ReportAllocs()
 	p := newParserWithGo()
-	arena.NewArena().Free()
 	b.ResetTimer()
 
-	parseAndInterpretInParallel(10, b.N, true, p)
+	parseAndInterpretInParallel(b, 10, b.N*10, true, p)
 }
 
 func BenchmarkParseAndInterpretWithoutArena(b *testing.B) {
 	b.ReportAllocs()
 	p := newParserWithGo()
-	arena.NewArena().Free()
 	b.ResetTimer()
 
-	parseAndInterpretInParallel(10, b.N, false, p)
+	parseAndInterpretInParallel(b,10, b.N*10, false, p)
 }
 
 func newParserWithGo() *Parser {
@@ -180,22 +181,25 @@ func newParserWithGo() *Parser {
 	return p
 }
 
-func parseAndInterpretInParallel(threads, repeats int, withArena bool, p *Parser) {
+func parseAndInterpretInParallel(b *testing.B, threads, repeats int, withArena bool, p *Parser) {
 	wg := new(sync.WaitGroup)
 	wg.Add(threads)
+	pool := heap.NewPool(threads, 10, time.Second)
 	for thread := 0; thread < threads; thread++ {
 		go func(thread int) {
 			for repeat := 0; repeat < repeats; repeat++ {
 				pkg := fmt.Sprintf("src/asp/parse_%v_%v", thread, repeat)
-				var heap *arena.Arena
+				var heap *heap.Heap
+				var arena *arena.Arena
 				if withArena {
-					heap = arena.NewArena()
+					heap = pool.Get()
+					arena = heap.Arena
 				}
 
-				s := p.interpreter.scope.newScope(core.NewPackage(pkg), heap, core.ParseModeNormal, filepath.Join(pkg, "BUILD"), 10)
+				s := p.interpreter.scope.newScope(core.NewPackage(pkg), arena, core.ParseModeNormal, filepath.Join(pkg, "BUILD"), 10)
 
 				// This call to ParseData currently doesn't use the arena
-				stmts, err := s.interpreter.parser.ParseData(heap, []byte(code), "BUILD")
+				stmts, err := s.interpreter.parser.ParseData(arena, []byte(code), "BUILD")
 				if err != nil {
 					panic(err)
 				}
@@ -206,12 +210,24 @@ func parseAndInterpretInParallel(threads, repeats int, withArena bool, p *Parser
 				t := s.state.Graph.TargetOrDie(core.NewBuildLabel(pkg, "asp"))
 				_ = t
 				if withArena {
-					heap.Free()
+					pool.Return(heap)
 				}
 			}
 			wg.Done()
 		}(thread)
 	}
 
 	wg.Wait()
+
+
+	gcStats := debug.GCStats{}
+	memStats := runtime.MemStats{}
+	debug.ReadGCStats(&gcStats)
+	runtime.ReadMemStats(&memStats)
+
+	b.ReportMetric(float64(pool.Stats.Frees.Load()), "Arena-frees")
+	b.ReportMetric(float64(pool.Stats.NewArena.Load()), "Arena-creates")
+	b.ReportMetric(float64(memStats.Sys)/1024.0/1024.0/1024.0, "In-use")
+	b.ReportMetric(float64(gcStats.NumGC), "GCs")
+
 }

src/parse/asp/heap/heap.go

@@ -1,9 +1,121 @@
+// Package heap provides a pool to acquire memory arenas to use to allocate objects when parsing and interpreting asp.
+// This package pools the areas to improve efficiency, as short-lived arenas, that allocate a handful of small objects
+// don't perform well. By re-using arenas between package parses, we can have them live longer and allocate more
+// objects.
 package heap
 
 import (
 	"arena"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/thought-machine/please/src/cli/logging"
 )
 
+var log = logging.Log
+
+// Heap represents a memory arena we can use for asp heap allocated objects
+type Heap struct {
+	usages    int
+	lastUsage time.Time
+	Arena     *arena.Arena
+	mux       sync.Mutex
+}
+
+func (h *Heap) free() {
+	if h.Arena == nil {
+		return
+	}
+	h.Arena.Free()
+	h.Arena = nil
+	h.usages = 0
+	h.lastUsage = time.Time{}
+}
+
+// Pool is a struct for managing a pool of heaps, freeing them after a set number of usages, or if they've not been used
+// for a set duration. This allows the heaps to be freed once we finish parsing.
+type Pool struct {
+	heaps     []*Heap
+	available chan *Heap
+	heapsMux  sync.Mutex
+	// Ideally this would be based on bytes allocated, but we don't have access to those stats.
+	UsagesBeforeFree  int
+	IdleTimeUntilFree time.Duration
+	Stats Stats
+}
+
+type Stats struct {
+	Frees atomic.Uint64
+	NewArena atomic.Uint64
+}
+
+// NewPool creates a new dynamically sized pool of heaps that can be used to allocated memory during parsing and
+// interpreting asp code.
+//
+// usagesBeforeFree:  The number of times a heap will be used before the pool frees the underlying arena. This can be
+//
+//	negative, in which case the arena will never be freed by this heuristic.
+//
+// idleTimeUntilFree: The duration in which the underlying arena will be freed if this heap is not used. Idle time is
+//
+//	calculated from when the heap was returned to the pool.
+func NewPool(size, usagesBeforeFree int, idleTimeUntilFree time.Duration) *Pool {
+	pool := &Pool{
+		heaps:             make([]*Heap, size),
+		UsagesBeforeFree:  usagesBeforeFree,
+		IdleTimeUntilFree: idleTimeUntilFree,
+		available:         make(chan *Heap, size),
+	}
+	for i := 0; i < size; i++ {
+		pool.heaps[i] = new(Heap)
+		pool.available <- pool.heaps[i]
+	}
+
+	go pool.freeIdleHeaps()
+
+	return pool
+}
+
+func (p *Pool) freeIdleHeaps() {
+	t := time.NewTicker(time.Second)
+	for {
+		select {
+		case <-t.C:
+			for _, h := range p.heaps {
+				if time.Since(h.lastUsage) < p.IdleTimeUntilFree {
+					continue
+				}
+				if !h.mux.TryLock() {
+					continue // Something must be using it so it's not idle
+				}
+				h.free()
+				h.mux.Unlock()
+			}
+		}
+	}
+}
+
+func (p *Pool) Get() *Heap {
+	var heap = <-p.available
+	heap.mux.Lock()
+	if heap.Arena == nil {
+		heap.Arena = arena.NewArena()
+		p.Stats.NewArena.Add(1)
+	}
+	return heap
+}
+
+func (p *Pool) Return(heap *Heap) {
+	heap.usages++
+	if heap.usages > p.UsagesBeforeFree {
+		heap.free()
+		p.Stats.Frees.Add(1)
+	}
+	heap.mux.Unlock()
+	p.available <- heap
+}
+
 func MakeSlice[T any](a *arena.Arena, len, cap int) []T {
 	if a == nil {
 		return make([]T, len, cap)