From b5a81ee1cb5524724aa818e0fdeafa54281806f8 Mon Sep 17 00:00:00 2001 From: John Hopper Date: Thu, 27 Aug 2026 10:26:06 -0700 Subject: [PATCH 1/3] feat(pg): add bounded translation cache --- drivers/pg/translation_cache.go | 407 ++++++++++++++++ .../pg/translation_cache_benchmark_test.go | 175 +++++++ drivers/pg/translation_cache_test.go | 450 ++++++++++++++++++ 3 files changed, 1032 insertions(+) create mode 100644 drivers/pg/translation_cache.go create mode 100644 drivers/pg/translation_cache_benchmark_test.go create mode 100644 drivers/pg/translation_cache_test.go diff --git a/drivers/pg/translation_cache.go b/drivers/pg/translation_cache.go new file mode 100644 index 00000000..ed8a8e9a --- /dev/null +++ b/drivers/pg/translation_cache.go @@ -0,0 +1,407 @@ +package pg + +import ( + "context" + "crypto/sha256" + "fmt" + "reflect" + "sort" + "strings" + "sync" + "sync/atomic" + + "github.com/specterops/dawgs/cache" + "github.com/specterops/dawgs/cypher/models/pgsql" +) + +const ( + translationCacheCapacity = 256 + translationCacheKeyFormat = 3 + maxCachedCypherBytes = 64 * 1024 + translationCachePolicy = "compiler-v4:optimized" +) + +type translationCacheKey struct { + format uint8 + queryDigest [sha256.Size]byte + querySize int + graphID int32 + parameterTypes string + policy string + schemaGeneration uint64 +} + +type translationCacheBuildResult struct { + parameters map[string]any + parameterSources map[string]string +} + +type translationCacheEntry struct { + sql string + parameterSources map[string]string +} + +type translationCacheCall struct { + done chan struct{} + doneOnce sync.Once + reusable bool +} + +func (s *translationCacheCall) finish(reusable bool) { + s.doneOnce.Do(func() { + s.reusable = reusable + close(s.done) + }) +} + +// TranslationCacheStats contains aggregate counters only. It intentionally +// exposes no source text, SQL, parameter names, values, or connection data. +type TranslationCacheStats struct { + Hits int64 + Misses int64 + Coalesced int64 + Bypasses int64 + Insertions int64 + Evictions int64 + BindingFailures int64 + BuildFailures int64 + UnoptimizedCompilations int64 + Size int64 + Capacity int + Generation uint64 +} + +type translationCache struct { + capacity int + policy string + entries cache.Cache[translationCacheKey, translationCacheEntry] + + lock sync.Mutex + pending map[translationCacheKey]*translationCacheCall + closed bool + + schemaGeneration atomic.Uint64 + hits atomic.Int64 + misses atomic.Int64 + coalesced atomic.Int64 + bypasses atomic.Int64 + insertions atomic.Int64 + evictions atomic.Int64 + bindingFailures atomic.Int64 + buildFailures atomic.Int64 + unoptimizedCompilations atomic.Int64 +} + +type translationCacheProvider interface { + TranslationCache() *translationCache +} + +type sharedTranslationCacheProvider struct { + cache *translationCache +} + +func (s sharedTranslationCacheProvider) TranslationCache() *translationCache { + return s.cache +} + +func newTranslationCache(capacity int) *translationCache { + return newTranslationCacheWithPolicy(capacity, translationCachePolicy) +} + +func newTranslationCacheWithPolicy(capacity int, policy string) *translationCache { + if capacity < 0 { + capacity = 0 + } + if policy == "" { + policy = translationCachePolicy + } + + cacheInstance := &translationCache{ + capacity: capacity, + policy: policy, + pending: map[translationCacheKey]*translationCacheCall{}, + } + if capacity > 0 { + cacheInstance.entries = cache.NewSieve[translationCacheKey, translationCacheEntry](capacity) + } + + return cacheInstance +} + +func (s *translationCache) Key(query string, graphID int32, parameters map[string]any) translationCacheKey { + query = strings.TrimSpace(query) + + return translationCacheKey{ + format: translationCacheKeyFormat, + queryDigest: sha256.Sum256([]byte(query)), + querySize: len(query), + graphID: graphID, + parameterTypes: parameterTypeShape(parameters), + policy: s.policy, + schemaGeneration: s.schemaGeneration.Load(), + } +} + +func parameterTypeShape(parameters map[string]any) string { + if len(parameters) == 0 { + return "" + } + + names := make([]string, 0, len(parameters)) + for name := range parameters { + names = append(names, name) + } + sort.Strings(names) + + shape := strings.Builder{} + for _, name := range names { + dataType, err := pgsql.ValueToDataType(parameters[name]) + if err != nil { + dataType = pgsql.DataType("unsupported:" + stableTypeName(parameters[name])) + } + + fmt.Fprintf(&shape, "%d:%s=%s;", len(name), name, dataType) + } + + return shape.String() +} + +func stableTypeName(value any) string { + if value == nil { + return "nil" + } + + return reflect.TypeOf(value).String() +} + +func (s *translationCache) GetOrBuild(key translationCacheKey, parameters map[string]any, build func() (string, translationCacheBuildResult, error)) (string, map[string]any, error) { + return s.GetOrBuildContext(context.Background(), key, parameters, build) +} + +func (s *translationCache) GetOrBuildContext(ctx context.Context, key translationCacheKey, parameters map[string]any, build func() (string, translationCacheBuildResult, error)) (string, map[string]any, error) { + if ctx == nil { + ctx = context.Background() + } + if key.querySize > maxCachedCypherBytes || s.capacity == 0 { + s.bypasses.Add(1) + return s.buildUncached(build) + } + + for { + if err := ctx.Err(); err != nil { + return "", nil, err + } + + s.lock.Lock() + if s.closed { + s.lock.Unlock() + s.bypasses.Add(1) + return s.buildUncached(build) + } + if s.entries != nil { + if entry, found := s.entries.Get(key); found { + s.lock.Unlock() + bindings, err := bindCachedParameters(entry.parameterSources, parameters) + if err != nil { + s.bindingFailures.Add(1) + return "", nil, err + } + s.hits.Add(1) + return entry.sql, bindings, nil + } + } + + if pending := s.pending[key]; pending != nil { + s.lock.Unlock() + select { + case <-pending.done: + s.coalesced.Add(1) + if !pending.reusable { + s.bypasses.Add(1) + return s.buildUncached(build) + } + continue + + case <-ctx.Done(): + return "", nil, ctx.Err() + } + } + + pending := &translationCacheCall{ + done: make(chan struct{}), + } + s.pending[key] = pending + s.misses.Add(1) + s.lock.Unlock() + + sql, result, err, panicked := callTranslationBuild(build) + entry, cacheable := cacheableTranslation(sql, result, parameters, err) + + s.lock.Lock() + currentGeneration := s.schemaGeneration.Load() + generationCurrent := key.schemaGeneration == currentGeneration + closed := s.closed + reusable := cacheable && generationCurrent && !closed && panicked == nil + if reusable { + if s.entries.Stats().Size() >= int64(s.capacity) { + s.evictions.Add(1) + } + s.entries.Put(cloneTranslationCacheKey(key), entry) + s.insertions.Add(1) + } + if s.pending[key] == pending { + delete(s.pending, key) + } + pending.finish(reusable) + s.lock.Unlock() + + if panicked != nil { + panic(panicked) + } + if err != nil { + s.buildFailures.Add(1) + return "", nil, err + } + if !generationCurrent && !closed { + key.schemaGeneration = currentGeneration + continue + } + if closed || !cacheable { + s.bypasses.Add(1) + return sql, result.parameters, nil + } + + return sql, result.parameters, nil + } +} + +func callTranslationBuild(build func() (string, translationCacheBuildResult, error)) (sql string, result translationCacheBuildResult, err error, panicked any) { + defer func() { + panicked = recover() + }() + + sql, result, err = build() + return sql, result, err, nil +} + +func (s *translationCache) buildUncached(build func() (string, translationCacheBuildResult, error)) (string, map[string]any, error) { + sql, result, err, panicked := callTranslationBuild(build) + if panicked != nil { + panic(panicked) + } + if err != nil { + s.buildFailures.Add(1) + return "", nil, err + } + + return sql, result.parameters, nil +} + +func (s *translationCache) BuildUnoptimized(build func() (string, translationCacheBuildResult, error)) (string, map[string]any, error) { + s.unoptimizedCompilations.Add(1) + s.bypasses.Add(1) + return s.buildUncached(build) +} + +func cloneTranslationCacheKey(key translationCacheKey) translationCacheKey { + key.parameterTypes = strings.Clone(key.parameterTypes) + key.policy = strings.Clone(key.policy) + return key +} + +func cacheableTranslation(sql string, result translationCacheBuildResult, parameters map[string]any, err error) (translationCacheEntry, bool) { + if err != nil { + return translationCacheEntry{}, false + } + if len(result.parameters) != len(result.parameterSources) { + return translationCacheEntry{}, false + } + + sources := make(map[string]string, len(result.parameterSources)) + for generated := range result.parameters { + source, found := result.parameterSources[generated] + if !found { + return translationCacheEntry{}, false + } + if source == "" { + return translationCacheEntry{}, false + } + if _, found := parameters[source]; !found { + return translationCacheEntry{}, false + } + sources[strings.Clone(generated)] = strings.Clone(source) + } + + return translationCacheEntry{ + sql: strings.Clone(sql), + parameterSources: sources, + }, true +} + +func bindCachedParameters(sources map[string]string, parameters map[string]any) (map[string]any, error) { + if len(sources) == 0 { + return nil, nil + } + + bindings := make(map[string]any, len(sources)) + for generated, source := range sources { + value, found := parameters[source] + if !found { + return nil, fmt.Errorf("missing value for Cypher parameter %q", source) + } + negotiated, err := pgsql.NegotiateValue(value) + if err != nil { + return nil, err + } + bindings[generated] = negotiated + } + + return bindings, nil +} + +func (s *translationCache) Invalidate() { + s.lock.Lock() + defer s.lock.Unlock() + + s.schemaGeneration.Add(1) + if !s.closed && s.capacity > 0 { + s.entries = cache.NewSieve[translationCacheKey, translationCacheEntry](s.capacity) + } +} + +func (s *translationCache) Close() { + s.lock.Lock() + defer s.lock.Unlock() + + s.closed = true + s.entries = nil + for key, pending := range s.pending { + delete(s.pending, key) + pending.finish(false) + } +} + +func (s *translationCache) Stats() TranslationCacheStats { + s.lock.Lock() + defer s.lock.Unlock() + + var size int64 + if s.entries != nil { + size = s.entries.Stats().Size() + } + + return TranslationCacheStats{ + Hits: s.hits.Load(), + Misses: s.misses.Load(), + Coalesced: s.coalesced.Load(), + Bypasses: s.bypasses.Load(), + Insertions: s.insertions.Load(), + Evictions: s.evictions.Load(), + BindingFailures: s.bindingFailures.Load(), + BuildFailures: s.buildFailures.Load(), + UnoptimizedCompilations: s.unoptimizedCompilations.Load(), + Size: size, + Capacity: s.capacity, + Generation: s.schemaGeneration.Load(), + } +} diff --git a/drivers/pg/translation_cache_benchmark_test.go b/drivers/pg/translation_cache_benchmark_test.go new file mode 100644 index 00000000..260ef284 --- /dev/null +++ b/drivers/pg/translation_cache_benchmark_test.go @@ -0,0 +1,175 @@ +package pg + +import ( + "context" + "testing" + + "github.com/specterops/dawgs/cypher/frontend" + "github.com/specterops/dawgs/cypher/models/pgsql/translate" + "github.com/specterops/dawgs/drivers/pg/pgutil" + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/query" +) + +func benchmarkTranslationBuild(query string, parameters map[string]any) func() (string, translationCacheBuildResult, error) { + return func() (string, translationCacheBuildResult, error) { + parsedQuery, err := frontend.ParseCypher(frontend.NewContext(), query) + if err != nil { + return "", translationCacheBuildResult{}, err + } + + kindMapper := pgutil.NewInMemoryKindMapper() + kindMapper.Put(graph.StringKind("NodeKind1")) + translation, parameterSources, err := translate.TranslateWithOptionsAndParameterSources(context.Background(), parsedQuery, kindMapper, parameters, translate.DefaultGraphID, translate.DefaultOptions()) + if err != nil { + return "", translationCacheBuildResult{}, err + } + + sql, err := translate.Translated(translation) + return sql, translationCacheBuildResult{ + parameters: translation.Parameters, + parameterSources: parameterSources, + }, err + } +} + +func BenchmarkTranslationCacheUncached(b *testing.B) { + build := benchmarkTranslationBuild(`MATCH (n:NodeKind1) WHERE n.name = $name RETURN n`, map[string]any{"name": "first"}) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, _, err := build(); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkTranslationCacheColdPopulation(b *testing.B) { + parameters := map[string]any{"name": "first"} + build := benchmarkTranslationBuild(`MATCH (n:NodeKind1) WHERE n.name = $name RETURN n`, parameters) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + translationCache := newTranslationCache(1) + if _, _, err := translationCache.GetOrBuild(translationCache.Key(`MATCH (n:NodeKind1) WHERE n.name = $name RETURN n`, 1, parameters), parameters, build); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkTranslationCacheParameterizedHit(b *testing.B) { + parameters := map[string]any{"name": "first"} + query := `MATCH (n:NodeKind1) WHERE n.name = $name RETURN n` + translationCache := newTranslationCache(1) + key := translationCache.Key(query, 1, parameters) + build := benchmarkTranslationBuild(query, parameters) + if _, _, err := translationCache.GetOrBuild(key, parameters, build); err != nil { + b.Fatal(err) + } + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + parameters["name"] = "next" + if _, _, err := translationCache.GetOrBuild(translationCache.Key(query, 1, parameters), parameters, build); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkTranslationCacheParameterlessHit(b *testing.B) { + query := `MATCH (n:NodeKind1) RETURN n` + translationCache := newTranslationCache(1) + key := translationCache.Key(query, 1, nil) + build := benchmarkTranslationBuild(query, nil) + if _, _, err := translationCache.GetOrBuild(key, nil, build); err != nil { + b.Fatal(err) + } + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, _, err := translationCache.GetOrBuild(key, nil, build); err != nil { + b.Fatal(err) + } + } +} + +func benchmarkPreparedFetchStartNodes(b *testing.B, id graph.ID) preparedRegularQuery { + b.Helper() + builder := query.NewBuilder(nil) + builder.Apply( + query.Where(query.Equals(query.StartProperty("objectid"), id)), + query.Returning(query.Relationship(), query.Start()), + ) + regularQuery, err := builder.Build(false) + if err != nil { + b.Fatal(err) + } + prepared, err := prepareRegularQuery(regularQuery) + if err != nil { + b.Fatal(err) + } + + return prepared +} + +func BenchmarkBuilderFetchStartNodesOptimizedCacheDisabled(b *testing.B) { + setOptimizedTranslationForTest(b, true) + manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 0, + }) + b.ReportAllocs() + b.ResetTimer() + for index := range b.N { + if _, _, err := manager.compileRegularQuery(context.Background(), benchmarkPreparedFetchStartNodes(b, graph.ID(index+1)), translate.DefaultGraphID); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkBuilderFetchStartNodesColdPopulation(b *testing.B) { + setOptimizedTranslationForTest(b, true) + b.ReportAllocs() + b.ResetTimer() + for index := range b.N { + manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 1, + }) + if _, _, err := manager.compileRegularQuery(context.Background(), benchmarkPreparedFetchStartNodes(b, graph.ID(index+1)), translate.DefaultGraphID); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkBuilderFetchStartNodesWarmHit(b *testing.B) { + setOptimizedTranslationForTest(b, true) + manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 1, + }) + if _, _, err := manager.compileRegularQuery(context.Background(), benchmarkPreparedFetchStartNodes(b, graph.ID(1)), translate.DefaultGraphID); err != nil { + b.Fatal(err) + } + + b.ReportAllocs() + b.ResetTimer() + for index := range b.N { + if _, _, err := manager.compileRegularQuery(context.Background(), benchmarkPreparedFetchStartNodes(b, graph.ID(index+2)), translate.DefaultGraphID); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkBuilderFetchStartNodesUnoptimized(b *testing.B) { + setOptimizedTranslationForTest(b, false) + manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 1, + }) + b.ReportAllocs() + b.ResetTimer() + for index := range b.N { + if _, _, err := manager.compileRegularQuery(context.Background(), benchmarkPreparedFetchStartNodes(b, graph.ID(index+1)), translate.DefaultGraphID); err != nil { + b.Fatal(err) + } + } +} diff --git a/drivers/pg/translation_cache_test.go b/drivers/pg/translation_cache_test.go new file mode 100644 index 00000000..356c15e7 --- /dev/null +++ b/drivers/pg/translation_cache_test.go @@ -0,0 +1,450 @@ +package pg + +import ( + "context" + "errors" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/specterops/dawgs/graph" + "github.com/stretchr/testify/require" +) + +func cacheableBuild(sql string, parameters map[string]any, sources map[string]string) func() (string, translationCacheBuildResult, error) { + return func() (string, translationCacheBuildResult, error) { + return sql, translationCacheBuildResult{ + parameters: parameters, + parameterSources: sources, + }, nil + } +} + +func TestTranslationCacheRebindsCurrentValues(t *testing.T) { + translationCache := newTranslationCache(2) + firstParameters := map[string]any{"needle": graph.ID(1)} + key := translationCache.Key("RETURN $needle", 7, firstParameters) + + _, bindings, err := translationCache.GetOrBuild(key, firstParameters, cacheableBuild("select @p0", map[string]any{"p0": uint64(1)}, map[string]string{"p0": "needle"})) + require.NoError(t, err) + require.Equal(t, map[string]any{"p0": uint64(1)}, bindings) + + secondParameters := map[string]any{"needle": graph.ID(2)} + builds := 0 + _, bindings, err = translationCache.GetOrBuild(translationCache.Key("RETURN $needle", 7, secondParameters), secondParameters, func() (string, translationCacheBuildResult, error) { + builds++ + return "", translationCacheBuildResult{}, nil + }) + require.NoError(t, err) + require.Zero(t, builds) + require.Equal(t, map[string]any{"p0": uint64(2)}, bindings) + require.Equal(t, int64(1), translationCache.Stats().Hits) +} + +func TestTranslationCachePartitionsKeysAndBypassesUnsafeResults(t *testing.T) { + translationCache := newTranslationCache(2) + parameters := map[string]any{"value": int64(1)} + key := translationCache.Key("RETURN $value", 1, parameters) + + require.NotEqual(t, key, translationCache.Key("RETURN $other", 1, parameters)) + require.NotEqual(t, key, translationCache.Key("RETURN $value", 2, parameters)) + require.NotEqual(t, key, translationCache.Key("RETURN $value", 1, map[string]any{"other": int64(1)})) + require.NotEqual(t, key, translationCache.Key("RETURN $value", 1, map[string]any{"value": "1"})) + translationCache.Invalidate() + require.NotEqual(t, key, translationCache.Key("RETURN $value", 1, parameters)) + key = translationCache.Key("RETURN $value", 1, parameters) + + builds := 0 + for range 2 { + _, _, err := translationCache.GetOrBuild(key, parameters, func() (string, translationCacheBuildResult, error) { + builds++ + return "select @p0", translationCacheBuildResult{ + parameters: map[string]any{"p0": int64(1)}, + }, nil + }) + require.NoError(t, err) + } + require.Equal(t, 2, builds, "a parameter without provenance must not be retained") +} + +func TestTranslationCacheKeyUsesDigestWithoutSourceText(t *testing.T) { + translationCache := newTranslationCache(1) + key := translationCache.Key("RETURN $secret // customer-value", 1, map[string]any{"secret": "value"}) + other := translationCache.Key("RETURN $other", 1, map[string]any{"secret": "value"}) + + require.Equal(t, len("RETURN $secret // customer-value"), key.querySize) + require.NotEqual(t, key.queryDigest, other.queryDigest) +} + +func TestTranslationCacheBypassesOversizedSource(t *testing.T) { + translationCache := newTranslationCache(1) + parameters := map[string]any{"value": int64(1)} + key := translationCache.Key(strings.Repeat("x", maxCachedCypherBytes+1), 1, parameters) + builds := 0 + + for range 2 { + _, _, err := translationCache.GetOrBuild(key, parameters, func() (string, translationCacheBuildResult, error) { + builds++ + return "select @p0", translationCacheBuildResult{ + parameters: map[string]any{"p0": int64(1)}, + parameterSources: map[string]string{"p0": "value"}, + }, nil + }) + require.NoError(t, err) + } + + require.Equal(t, 2, builds) + require.Equal(t, int64(2), translationCache.Stats().Bypasses) +} + +func TestCacheableTranslationRequiresExactParameterProvenance(t *testing.T) { + _, cacheable := cacheableTranslation("select @p0", translationCacheBuildResult{ + parameters: map[string]any{"p0": int64(1)}, + parameterSources: map[string]string{"p1": "value"}, + }, map[string]any{"value": int64(1)}, nil) + + require.False(t, cacheable) +} + +func TestTranslationCacheReportsEviction(t *testing.T) { + translationCache := newTranslationCache(1) + parameters := map[string]any{"value": int64(1)} + + _, _, err := translationCache.GetOrBuild(translationCache.Key("RETURN $value", 1, parameters), parameters, cacheableBuild("select @p0", map[string]any{"p0": int64(1)}, map[string]string{"p0": "value"})) + require.NoError(t, err) + _, _, err = translationCache.GetOrBuild(translationCache.Key("RETURN $value + 1", 1, parameters), parameters, cacheableBuild("select @p0", map[string]any{"p0": int64(1)}, map[string]string{"p0": "value"})) + require.NoError(t, err) + + stats := translationCache.Stats() + require.Equal(t, int64(1), stats.Evictions) + require.Equal(t, int64(1), stats.Size) +} + +func TestTranslationCacheCoalescesCacheableMisses(t *testing.T) { + translationCache := newTranslationCache(2) + parameters := map[string]any{"value": int64(1)} + key := translationCache.Key("RETURN $value", 1, parameters) + started := make(chan struct{}) + release := make(chan struct{}) + errs := make(chan error, 8) + var builds atomic.Int64 + var group sync.WaitGroup + + for range 8 { + group.Add(1) + go func() { + defer group.Done() + _, _, err := translationCache.GetOrBuild(key, parameters, func() (string, translationCacheBuildResult, error) { + if builds.Add(1) == 1 { + close(started) + <-release + } + return "select @p0", translationCacheBuildResult{ + parameters: map[string]any{"p0": int64(1)}, + parameterSources: map[string]string{"p0": "value"}, + }, nil + }) + errs <- err + }() + } + + <-started + time.Sleep(10 * time.Millisecond) + close(release) + group.Wait() + close(errs) + for err := range errs { + require.NoError(t, err) + } + require.Equal(t, int64(1), builds.Load()) +} + +func TestTranslationCacheCoalescedRequestsBindTheirOwnValues(t *testing.T) { + translationCache := newTranslationCache(2) + source := "RETURN $value" + leaderParameters := map[string]any{"value": graph.ID(1)} + key := translationCache.Key(source, 1, leaderParameters) + started := make(chan struct{}) + release := make(chan struct{}) + + type result struct { + bindings map[string]any + err error + } + leaderDone := make(chan result, 1) + go func() { + _, bindings, err := translationCache.GetOrBuild(key, leaderParameters, func() (string, translationCacheBuildResult, error) { + close(started) + <-release + return "select @p0", translationCacheBuildResult{ + parameters: map[string]any{"p0": uint64(1)}, + parameterSources: map[string]string{"p0": "value"}, + }, nil + }) + leaderDone <- result{ + bindings: bindings, + err: err, + } + }() + + <-started + values := []graph.ID{2, 3, 4, 5} + waiters := make(chan result, len(values)) + for _, value := range values { + value := value + go func() { + parameters := map[string]any{"value": value} + _, bindings, err := translationCache.GetOrBuild(translationCache.Key(source, 1, parameters), parameters, func() (string, translationCacheBuildResult, error) { + return "", translationCacheBuildResult{}, errors.New("coalesced request unexpectedly rebuilt the translation") + }) + waiters <- result{ + bindings: bindings, + err: err, + } + }() + } + + close(release) + leader := <-leaderDone + require.NoError(t, leader.err) + require.Equal(t, map[string]any{"p0": uint64(1)}, leader.bindings) + + seen := map[uint64]struct{}{} + for range values { + waiter := <-waiters + require.NoError(t, waiter.err) + value, ok := waiter.bindings["p0"].(uint64) + require.True(t, ok) + seen[value] = struct{}{} + } + require.Equal(t, map[uint64]struct{}{2: {}, 3: {}, 4: {}, 5: {}}, seen) +} + +func TestTranslationCacheWaiterCancellationDoesNotInterruptLeader(t *testing.T) { + translationCache := newTranslationCache(2) + parameters := map[string]any{"value": graph.ID(1)} + key := translationCache.Key("RETURN $value", 1, parameters) + started := make(chan struct{}) + release := make(chan struct{}) + leaderDone := make(chan error, 1) + + go func() { + _, _, err := translationCache.GetOrBuild(key, parameters, func() (string, translationCacheBuildResult, error) { + close(started) + <-release + return "select @p0", translationCacheBuildResult{ + parameters: map[string]any{"p0": uint64(1)}, + parameterSources: map[string]string{"p0": "value"}, + }, nil + }) + leaderDone <- err + }() + + <-started + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, _, err := translationCache.GetOrBuildContext(ctx, key, parameters, func() (string, translationCacheBuildResult, error) { + return "", translationCacheBuildResult{}, errors.New("cancelled request unexpectedly rebuilt the translation") + }) + require.ErrorIs(t, err, context.Canceled) + + close(release) + require.NoError(t, <-leaderDone) + _, bindings, err := translationCache.GetOrBuild(key, map[string]any{"value": graph.ID(2)}, func() (string, translationCacheBuildResult, error) { + return "", translationCacheBuildResult{}, errors.New("completed leader did not publish a reusable translation") + }) + require.NoError(t, err) + require.Equal(t, map[string]any{"p0": uint64(2)}, bindings) +} + +func TestTranslationCacheParameterlessHitAndClose(t *testing.T) { + translationCache := newTranslationCache(1) + key := translationCache.Key("RETURN 1", 1, nil) + build := cacheableBuild("select 1", map[string]any{}, map[string]string{}) + + _, bindings, err := translationCache.GetOrBuild(key, nil, build) + require.NoError(t, err) + require.Empty(t, bindings) + _, bindings, err = translationCache.GetOrBuild(key, nil, build) + require.NoError(t, err) + require.Nil(t, bindings) + + translationCache.Close() + builds := 0 + _, _, err = translationCache.GetOrBuild(key, nil, func() (string, translationCacheBuildResult, error) { + builds++ + return "select 1", translationCacheBuildResult{}, nil + }) + require.NoError(t, err) + require.Equal(t, 1, builds) +} + +func TestTranslationCacheNonCacheableWaitersBuildIndependently(t *testing.T) { + translationCache := newTranslationCache(2) + parameters := map[string]any{"value": int64(1)} + key := translationCache.Key("RETURN $value", 1, parameters) + leaderStarted := make(chan struct{}) + releaseLeader := make(chan struct{}) + uncachedStarted := make(chan struct{}, 7) + releaseUncached := make(chan struct{}) + errs := make(chan error, 8) + var builds atomic.Int64 + var group sync.WaitGroup + + for range 8 { + group.Add(1) + go func() { + defer group.Done() + _, _, err := translationCache.GetOrBuild(key, parameters, func() (string, translationCacheBuildResult, error) { + if builds.Add(1) == 1 { + close(leaderStarted) + <-releaseLeader + } else { + uncachedStarted <- struct{}{} + <-releaseUncached + } + + return "select @p0", translationCacheBuildResult{ + parameters: map[string]any{"p0": int64(1)}, + }, nil + }) + errs <- err + }() + } + + <-leaderStarted + time.Sleep(10 * time.Millisecond) + close(releaseLeader) + for range 7 { + select { + case <-uncachedStarted: + case <-time.After(time.Second): + t.Fatal("non-cacheable waiters were serialized") + } + } + close(releaseUncached) + group.Wait() + close(errs) + for err := range errs { + require.NoError(t, err) + } + require.Equal(t, int64(8), builds.Load()) +} + +func TestTranslationCachePanicReleasesWaiters(t *testing.T) { + translationCache := newTranslationCache(2) + parameters := map[string]any{"value": int64(1)} + key := translationCache.Key("RETURN $value", 1, parameters) + leaderStarted := make(chan struct{}) + releaseLeader := make(chan struct{}) + leaderDone := make(chan struct{}) + panicValue := make(chan any, 1) + waiterDone := make(chan error, 1) + + go func() { + defer close(leaderDone) + defer func() { + panicValue <- recover() + }() + _, _, _ = translationCache.GetOrBuild(key, parameters, func() (string, translationCacheBuildResult, error) { + close(leaderStarted) + <-releaseLeader + panic("boom") + }) + }() + + <-leaderStarted + go func() { + _, _, err := translationCache.GetOrBuild(key, parameters, cacheableBuild( + "select @p0", + map[string]any{"p0": int64(1)}, + map[string]string{"p0": "value"}, + )) + waiterDone <- err + }() + + time.Sleep(10 * time.Millisecond) + close(releaseLeader) + <-leaderDone + require.Equal(t, "boom", <-panicValue) + select { + case err := <-waiterDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("panic left a cache waiter blocked") + } +} + +func TestTranslationCacheCloseReleasesWaiters(t *testing.T) { + translationCache := newTranslationCache(2) + parameters := map[string]any{"value": int64(1)} + key := translationCache.Key("RETURN $value", 1, parameters) + leaderStarted := make(chan struct{}) + releaseLeader := make(chan struct{}) + leaderDone := make(chan struct{}) + waiterDone := make(chan error, 1) + + go func() { + defer close(leaderDone) + _, _, _ = translationCache.GetOrBuild(key, parameters, func() (string, translationCacheBuildResult, error) { + close(leaderStarted) + <-releaseLeader + return "select @p0", translationCacheBuildResult{}, nil + }) + }() + + <-leaderStarted + go func() { + _, _, err := translationCache.GetOrBuildContext(context.Background(), key, parameters, cacheableBuild( + "select @p0", + map[string]any{"p0": int64(1)}, + map[string]string{"p0": "value"}, + )) + waiterDone <- err + }() + + time.Sleep(10 * time.Millisecond) + translationCache.Close() + select { + case err := <-waiterDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("cache close left a waiter blocked") + } + close(releaseLeader) + <-leaderDone +} + +func TestTranslationCacheInvalidationDoesNotPublishStaleBuild(t *testing.T) { + translationCache := newTranslationCache(2) + parameters := map[string]any{"value": int64(1)} + key := translationCache.Key("RETURN $value", 1, parameters) + started := make(chan struct{}) + release := make(chan struct{}) + var builds atomic.Int64 + + done := make(chan error, 1) + go func() { + _, _, err := translationCache.GetOrBuild(key, parameters, func() (string, translationCacheBuildResult, error) { + if builds.Add(1) == 1 { + close(started) + <-release + } + return "select @p0", translationCacheBuildResult{ + parameters: map[string]any{"p0": int64(1)}, + parameterSources: map[string]string{"p0": "value"}, + }, nil + }) + done <- err + }() + + <-started + translationCache.Invalidate() + close(release) + require.NoError(t, <-done) + require.Equal(t, int64(2), builds.Load()) + require.Equal(t, int64(1), translationCache.Stats().Size) + require.Equal(t, uint64(1), translationCache.Stats().Generation) +} From f059ea5d1dedeb46ba87316cc33ff1f0804a8e64 Mon Sep 17 00:00:00 2001 From: John Hopper Date: Thu, 27 Aug 2026 10:26:59 -0700 Subject: [PATCH 2/3] feat(pg): cache raw Cypher translations --- drivers/pg/compiler.go | 38 ++++++++++++++++++++++++++++++ drivers/pg/driver.go | 43 ++++++++++++++++++++++++++++++++-- drivers/pg/manager.go | 49 +++++++++++++++++++++++++++++++-------- drivers/pg/transaction.go | 17 ++++++-------- 4 files changed, 125 insertions(+), 22 deletions(-) create mode 100644 drivers/pg/compiler.go diff --git a/drivers/pg/compiler.go b/drivers/pg/compiler.go new file mode 100644 index 00000000..a971684d --- /dev/null +++ b/drivers/pg/compiler.go @@ -0,0 +1,38 @@ +package pg + +import ( + "context" + "strings" + + "github.com/specterops/dawgs/cypher/frontend" + "github.com/specterops/dawgs/cypher/models/cypher" + "github.com/specterops/dawgs/cypher/models/pgsql/translate" +) + +func (s *SchemaManager) compileText(ctx context.Context, source string, parameters map[string]any, graphID int32) (string, map[string]any, error) { + return s.compile(ctx, strings.TrimSpace(source), parameters, graphID, func() (*cypher.RegularQuery, error) { + return frontend.ParseCypher(frontend.NewContext(), source) + }) +} + +func (s *SchemaManager) compile(ctx context.Context, source string, parameters map[string]any, graphID int32, parse func() (*cypher.RegularQuery, error)) (string, map[string]any, error) { + translationCache := s.translationCacheProvider.TranslationCache() + + build := func() (string, translationCacheBuildResult, error) { + if regularQuery, err := parse(); err != nil { + return "", translationCacheBuildResult{}, err + } else if translated, parameterSources, err := translate.TranslateWithOptionsAndParameterSources(ctx, regularQuery, s, parameters, graphID, translate.DefaultOptions()); err != nil { + return "", translationCacheBuildResult{}, err + } else if sqlQuery, err := translate.Translated(translated); err != nil { + return "", translationCacheBuildResult{}, err + } else { + return sqlQuery, translationCacheBuildResult{ + parameters: translated.Parameters, + parameterSources: parameterSources, + }, nil + } + } + + key := translationCache.Key(source, graphID, parameters) + return translationCache.GetOrBuildContext(ctx, key, parameters, build) +} diff --git a/drivers/pg/driver.go b/drivers/pg/driver.go index 5b14e459..8df09a13 100644 --- a/drivers/pg/driver.go +++ b/drivers/pg/driver.go @@ -42,11 +42,38 @@ type Driver struct { *SchemaManager } +// DriverOptions configures driver-wide behavior. TranslationCacheEntries is a +// count, not a per-connection setting; zero disables translation caching. +type DriverOptions struct { + TranslationCacheEntries int +} + +// DefaultDriverOptions returns the driver-wide PostgreSQL settings used by +// NewDriver. +func DefaultDriverOptions() DriverOptions { + return DriverOptions{ + TranslationCacheEntries: translationCacheCapacity, + } +} + func NewDriver(graphQueryMemoryLimit size.Size, pool *pgxpool.Pool) *Driver { + return NewDriverWithOptions(graphQueryMemoryLimit, pool, DefaultDriverOptions()) +} + +// NewDriverWithOptions constructs a PostgreSQL driver with driver-wide options. +func NewDriverWithOptions(graphQueryMemoryLimit size.Size, pool *pgxpool.Pool, options DriverOptions) *Driver { + options = normalizeDriverOptions(options) return &Driver{ pool: pool, - SchemaManager: NewSchemaManager(pool, graphQueryMemoryLimit), + SchemaManager: NewSchemaManagerWithOptions(pool, graphQueryMemoryLimit, options), + } +} + +func normalizeDriverOptions(options DriverOptions) DriverOptions { + if options.TranslationCacheEntries < 0 { + options.TranslationCacheEntries = 0 } + return options } func (s *Driver) SetDefaultGraph(ctx context.Context, graphSchema graph.Graph) error { @@ -96,10 +123,17 @@ func (s *Driver) BatchOperation(ctx context.Context, batchDelegate graph.BatchDe } func (s *Driver) Close(ctx context.Context) error { + s.translationCache.Close() s.pool.Close() return nil } +// TranslationCacheStats returns aggregate PostgreSQL translation-cache +// counters. It never exposes cached query text, SQL, or parameter data. +func (s *Driver) TranslationCacheStats() TranslationCacheStats { + return s.translationCache.Stats() +} + func renderConfig(batchWriteSize int, pgxOptions pgx.TxOptions, userOptions []graph.TransactionOption) (*Config, error) { graphCfg := graph.TransactionConfig{ DriverConfig: &Config{ @@ -175,7 +209,12 @@ func (s *Driver) RefreshKinds(ctx context.Context) error { // Wipe this map to be rebuilt in the fetch call below s.SchemaManager.kindIDsByKind = map[int16]graph.Kind{} - return s.SchemaManager.Fetch(ctx) + if err := s.SchemaManager.Fetch(ctx); err != nil { + return err + } + + s.translationCache.Invalidate() + return nil } func (s *Driver) OptimizeStorage(ctx context.Context) error { diff --git a/drivers/pg/manager.go b/drivers/pg/manager.go index 4ce56419..cbd5a0c2 100644 --- a/drivers/pg/manager.go +++ b/drivers/pg/manager.go @@ -33,17 +33,37 @@ func KindMapperFromGraphDatabase(graphDB graph.Database) (KindMapper, error) { } type SchemaManager struct { - defaultGraph model.Graph - pool *pgxpool.Pool - hasDefaultGraph bool - graphs map[string]model.Graph - kindsByID map[graph.Kind]int16 - kindIDsByKind map[int16]graph.Kind - lock *sync.RWMutex - graphQueryMemoryLimit size.Size + defaultGraph model.Graph + pool *pgxpool.Pool + hasDefaultGraph bool + graphs map[string]model.Graph + kindsByID map[graph.Kind]int16 + kindIDsByKind map[int16]graph.Kind + lock *sync.RWMutex + graphQueryMemoryLimit size.Size + translationCache *translationCache + translationCacheProvider translationCacheProvider } func NewSchemaManager(pool *pgxpool.Pool, graphQueryMemoryLimit size.Size) *SchemaManager { + return NewSchemaManagerWithOptions(pool, graphQueryMemoryLimit, DefaultDriverOptions()) +} + +// NewSchemaManagerWithTranslationCache permits an application to disable the +// compilation cache (zero), or to choose a bounded driver-wide entry capacity. +// Negative values disable the cache. +func NewSchemaManagerWithTranslationCache(pool *pgxpool.Pool, graphQueryMemoryLimit size.Size, translationCacheEntries int) *SchemaManager { + return NewSchemaManagerWithOptions(pool, graphQueryMemoryLimit, DriverOptions{ + TranslationCacheEntries: translationCacheEntries, + }) +} + +// NewSchemaManagerWithOptions constructs the shared compilation service with +// a bounded translation cache. +func NewSchemaManagerWithOptions(pool *pgxpool.Pool, graphQueryMemoryLimit size.Size, options DriverOptions) *SchemaManager { + options = normalizeDriverOptions(options) + translationCache := newTranslationCache(options.TranslationCacheEntries) + return &SchemaManager{ pool: pool, hasDefaultGraph: false, @@ -52,6 +72,10 @@ func NewSchemaManager(pool *pgxpool.Pool, graphQueryMemoryLimit size.Size) *Sche kindIDsByKind: map[int16]graph.Kind{}, lock: &sync.RWMutex{}, graphQueryMemoryLimit: graphQueryMemoryLimit, + translationCache: translationCache, + translationCacheProvider: sharedTranslationCacheProvider{ + cache: translationCache, + }, } } @@ -406,7 +430,12 @@ func (s *SchemaManager) AssertSchema(ctx context.Context, schema graph.Schema) e s.lock.Lock() defer s.lock.Unlock() - return s.WriteTransaction(ctx, func(tx graph.Transaction) error { + if err := s.WriteTransaction(ctx, func(tx graph.Transaction) error { return s.assertSchema(tx, schema) - }, OptionSetQueryExecMode(pgx.QueryExecModeSimpleProtocol)) + }, OptionSetQueryExecMode(pgx.QueryExecModeSimpleProtocol)); err != nil { + return err + } + + s.translationCache.Invalidate() + return nil } diff --git a/drivers/pg/transaction.go b/drivers/pg/transaction.go index 7bf4bbd7..c99c06a3 100644 --- a/drivers/pg/transaction.go +++ b/drivers/pg/transaction.go @@ -5,12 +5,10 @@ import ( "fmt" "github.com/specterops/dawgs/cypher/models/pgsql" - "github.com/specterops/dawgs/cypher/models/pgsql/translate" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgxpool" - "github.com/specterops/dawgs/cypher/frontend" "github.com/specterops/dawgs/drivers/pg/model" "github.com/specterops/dawgs/graph" "github.com/specterops/dawgs/query" @@ -275,16 +273,15 @@ func (s *transaction) query(query string, parameters map[string]any) (pgx.Rows, } func (s *transaction) Query(query string, parameters map[string]any) graph.Result { - if parsedQuery, err := frontend.ParseCypher(frontend.NewContext(), query); err != nil { - return graph.NewErrorResult(err) - } else if graphTarget, err := s.getTargetGraph(); err != nil { - return graph.NewErrorResult(err) - } else if translated, err := translate.Translate(s.ctx, parsedQuery, s.schemaManager, parameters, graphTarget.ID); err != nil { - return graph.NewErrorResult(err) - } else if sqlQuery, err := translate.Translated(translated); err != nil { + if graphTarget, err := s.getTargetGraph(); err != nil { return graph.NewErrorResult(err) } else { - return s.Raw(sqlQuery, translated.Parameters) + sqlQuery, bindings, err := s.schemaManager.compileText(s.ctx, query, parameters, graphTarget.ID) + if err != nil { + return graph.NewErrorResult(err) + } + + return s.Raw(sqlQuery, bindings) } } From 0f159576ab997a90da3da0b2c0c2a24cb55a7be7 Mon Sep 17 00:00:00 2001 From: John Hopper Date: Mon, 31 Aug 2026 07:55:49 -0700 Subject: [PATCH 3/3] chore: test fixes from stack rebase --- bdd/utils_test.go | 2 +- .../pg/translation_cache_benchmark_test.go | 80 ------------------- 2 files changed, 1 insertion(+), 81 deletions(-) diff --git a/bdd/utils_test.go b/bdd/utils_test.go index ffeda0df..c68a59b4 100644 --- a/bdd/utils_test.go +++ b/bdd/utils_test.go @@ -21,7 +21,7 @@ import ( "testing" "github.com/cucumber/godog" - "github.com/cucumber/messages/go/v34" + messages "github.com/cucumber/messages/go/v34" "github.com/specterops/dawgs/graph" "github.com/stretchr/testify/require" ) diff --git a/drivers/pg/translation_cache_benchmark_test.go b/drivers/pg/translation_cache_benchmark_test.go index 260ef284..ca76a486 100644 --- a/drivers/pg/translation_cache_benchmark_test.go +++ b/drivers/pg/translation_cache_benchmark_test.go @@ -8,7 +8,6 @@ import ( "github.com/specterops/dawgs/cypher/models/pgsql/translate" "github.com/specterops/dawgs/drivers/pg/pgutil" "github.com/specterops/dawgs/graph" - "github.com/specterops/dawgs/query" ) func benchmarkTranslationBuild(query string, parameters map[string]any) func() (string, translationCacheBuildResult, error) { @@ -94,82 +93,3 @@ func BenchmarkTranslationCacheParameterlessHit(b *testing.B) { } } } - -func benchmarkPreparedFetchStartNodes(b *testing.B, id graph.ID) preparedRegularQuery { - b.Helper() - builder := query.NewBuilder(nil) - builder.Apply( - query.Where(query.Equals(query.StartProperty("objectid"), id)), - query.Returning(query.Relationship(), query.Start()), - ) - regularQuery, err := builder.Build(false) - if err != nil { - b.Fatal(err) - } - prepared, err := prepareRegularQuery(regularQuery) - if err != nil { - b.Fatal(err) - } - - return prepared -} - -func BenchmarkBuilderFetchStartNodesOptimizedCacheDisabled(b *testing.B) { - setOptimizedTranslationForTest(b, true) - manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ - TranslationCacheEntries: 0, - }) - b.ReportAllocs() - b.ResetTimer() - for index := range b.N { - if _, _, err := manager.compileRegularQuery(context.Background(), benchmarkPreparedFetchStartNodes(b, graph.ID(index+1)), translate.DefaultGraphID); err != nil { - b.Fatal(err) - } - } -} - -func BenchmarkBuilderFetchStartNodesColdPopulation(b *testing.B) { - setOptimizedTranslationForTest(b, true) - b.ReportAllocs() - b.ResetTimer() - for index := range b.N { - manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ - TranslationCacheEntries: 1, - }) - if _, _, err := manager.compileRegularQuery(context.Background(), benchmarkPreparedFetchStartNodes(b, graph.ID(index+1)), translate.DefaultGraphID); err != nil { - b.Fatal(err) - } - } -} - -func BenchmarkBuilderFetchStartNodesWarmHit(b *testing.B) { - setOptimizedTranslationForTest(b, true) - manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ - TranslationCacheEntries: 1, - }) - if _, _, err := manager.compileRegularQuery(context.Background(), benchmarkPreparedFetchStartNodes(b, graph.ID(1)), translate.DefaultGraphID); err != nil { - b.Fatal(err) - } - - b.ReportAllocs() - b.ResetTimer() - for index := range b.N { - if _, _, err := manager.compileRegularQuery(context.Background(), benchmarkPreparedFetchStartNodes(b, graph.ID(index+2)), translate.DefaultGraphID); err != nil { - b.Fatal(err) - } - } -} - -func BenchmarkBuilderFetchStartNodesUnoptimized(b *testing.B) { - setOptimizedTranslationForTest(b, false) - manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ - TranslationCacheEntries: 1, - }) - b.ReportAllocs() - b.ResetTimer() - for index := range b.N { - if _, _, err := manager.compileRegularQuery(context.Background(), benchmarkPreparedFetchStartNodes(b, graph.ID(index+1)), translate.DefaultGraphID); err != nil { - b.Fatal(err) - } - } -}