diff --git a/packages/go/analysis/ad/owns.go b/packages/go/analysis/ad/owns.go index 4a03bc88047a..161916efd7f1 100644 --- a/packages/go/analysis/ad/owns.go +++ b/packages/go/analysis/ad/owns.go @@ -18,6 +18,7 @@ package ad import ( "context" + "errors" "fmt" "log/slog" @@ -76,17 +77,17 @@ func PostOwnsAndWriteOwner(ctx context.Context, db graph.Database, localGroupDat } sink := post.NewFilteredRelationshipSink(ctx, "PostOwnsAndWriteOwner", db, ownsWriteOwnerTracker) - defer sink.Done() + defer sink.Abort() if err := postOwnsEdges(ctx, db, sink, dsHeuristicsCache, anyEnforced, adminGroupIds); err != nil { - return sink.Stats(), err + return sink.Stats(), errors.Join(err, sink.Abort()) } if err := postWriteOwnerEdges(ctx, db, sink, dsHeuristicsCache, anyEnforced); err != nil { - return sink.Stats(), err + return sink.Stats(), errors.Join(err, sink.Abort()) } - return sink.Stats(), nil + return sink.Stats(), sink.Done() } func postOwnsEdges(ctx context.Context, db graph.Database, sink *post.FilteredRelationshipSink, dsHeuristicsCache map[string]bool, anyEnforced bool, adminGroupIds cardinality.Duplex[uint64]) error { diff --git a/packages/go/analysis/azure/post.go b/packages/go/analysis/azure/post.go index 1a9f31a076c4..a4b162ccf3e2 100644 --- a/packages/go/analysis/azure/post.go +++ b/packages/go/analysis/azure/post.go @@ -18,6 +18,7 @@ package azure import ( "context" + "errors" "fmt" "log/slog" "strings" @@ -225,15 +226,15 @@ func AppRoleAssignments(ctx context.Context, db graph.Database) (*post.AtomicPos return &post.AtomicPostProcessingStats{}, err } else { sink := post.NewFilteredRelationshipSink(ctx, "Azure App Role Assignments Post Processing", db, appRoleAssignmentTracker) - defer sink.Done() + defer sink.Abort() for _, tenant := range tenants { if err := postAppRoleAssignmentsForTenant(ctx, db, sink, tenant); err != nil { - return &post.AtomicPostProcessingStats{}, err + return &post.AtomicPostProcessingStats{}, errors.Join(err, sink.Abort()) } } - return sink.Stats(), nil + return sink.Stats(), sink.Done() } } @@ -523,15 +524,15 @@ func CreateAZAddOwnerEdge(ctx context.Context, db graph.Database) (*post.AtomicP return &post.AtomicPostProcessingStats{}, err } else { sink := post.NewFilteredRelationshipSink(ctx, "Azure Add Owner Post Processing", db, addOwnerTracker) - defer sink.Done() + defer sink.Abort() for _, tenant := range tenantNodes { if err := postAzureAddOwner(ctx, db, sink, tenant); err != nil { - return &post.AtomicPostProcessingStats{}, err + return &post.AtomicPostProcessingStats{}, errors.Join(err, sink.Abort()) } } - return sink.Stats(), nil + return sink.Stats(), sink.Done() } } @@ -837,14 +838,14 @@ func UserRoleAssignments(ctx context.Context, db graph.Database) (*post.AtomicPo return &post.AtomicPostProcessingStats{}, err } else { sink := post.NewFilteredRelationshipSink(ctx, "Azure User Role Assignments Post Processing", db, userRoleAssignmentTracker) - defer sink.Done() + defer sink.Abort() for _, tenant := range tenantNodes { if roleAssignments, err := FetchTenantRoleAssignments(ctx, db, tenant); err != nil { - return &post.AtomicPostProcessingStats{}, err + return &post.AtomicPostProcessingStats{}, errors.Join(err, sink.Abort()) } else { if err := postAzureResetPassword(ctx, db, sink, tenant, roleAssignments); err != nil { - return &post.AtomicPostProcessingStats{}, err + return &post.AtomicPostProcessingStats{}, errors.Join(err, sink.Abort()) } else { postAzureGlobalAdmins(ctx, sink, roleAssignments, tenant) postAzurePrivilegedRoleAdmins(ctx, sink, roleAssignments, tenant) @@ -852,12 +853,13 @@ func UserRoleAssignments(ctx context.Context, db graph.Database) (*post.AtomicPo if err := postAzureAddMembers(ctx, sink, roleAssignments); err != nil { slog.Error("Azure AddMember Post-Processing Failure", attr.Error(err)) + return &post.AtomicPostProcessingStats{}, errors.Join(err, sink.Abort()) } } } } - return sink.Stats(), nil + return sink.Stats(), sink.Done() } } diff --git a/packages/go/analysis/post/sink.go b/packages/go/analysis/post/sink.go index d5eea0d2d29b..1291faeafd4c 100644 --- a/packages/go/analysis/post/sink.go +++ b/packages/go/analysis/post/sink.go @@ -18,6 +18,8 @@ package post import ( "context" + "errors" + "fmt" "log/slog" "runtime" "sync" @@ -38,8 +40,8 @@ func newPropertiesWithFirstSeen() *graph.Properties { // FilteredRelationshipSink is an asynchronous graph relationship writer that ensures only new relationships are // inserted and removes unused ones. It uses a delta tracker to track changes between graphs and avoids reinserting -// existing edges. Any edge not visited during processing is treated as obsolete and will be deleted after the -// operation completes. +// existing edges. Any edge not visited during successful processing is treated as obsolete and will be deleted +// by Done. Failed processing must call Abort to preserve existing edges. type FilteredRelationshipSink struct { operationName string db graph.Database @@ -47,19 +49,33 @@ type FilteredRelationshipSink struct { jobC chan EnsureRelationshipJob stats AtomicPostProcessingStats wg sync.WaitGroup + ctx context.Context + cancel context.CancelFunc + submitLock sync.RWMutex + submitDone chan struct{} + shutdownOnce sync.Once + errorLock sync.Mutex + workerError error + doneError error } // NewFilteredRelationshipSink creates a new filtered relationship sink initialized with a given database, delta tracker, and operation name. func NewFilteredRelationshipSink(ctx context.Context, operationName string, db graph.Database, deltaSubgraph *Tracker) *FilteredRelationshipSink { - newSink := &FilteredRelationshipSink{ - db: db, - edgeTracker: deltaSubgraph, - operationName: operationName, - jobC: make(chan EnsureRelationshipJob), - stats: NewAtomicPostProcessingStats(), - } + var ( + workerContext, cancel = context.WithCancel(ctx) + newSink = &FilteredRelationshipSink{ + db: db, + edgeTracker: deltaSubgraph, + operationName: operationName, + jobC: make(chan EnsureRelationshipJob), + stats: NewAtomicPostProcessingStats(), + ctx: workerContext, + cancel: cancel, + submitDone: make(chan struct{}), + } + ) - newSink.start(ctx) + newSink.start(workerContext) return newSink } @@ -82,7 +98,8 @@ func (s *FilteredRelationshipSink) insertWorker(ctx context.Context, commonProps } if err := batch.CreateRelationshipByIDs(nextJob.FromID, nextJob.ToID, nextJob.Kind, relProps); err != nil { - slog.Error("Create Relationship Error", slog.String("err", err.Error())) + s.recordError(fmt.Errorf("%s: creating relationship: %w", s.operationName, err)) + return err } else { s.stats.AddRelationshipsCreated(nextJob.Kind, 1) @@ -94,12 +111,27 @@ func (s *FilteredRelationshipSink) insertWorker(ctx context.Context, commonProps } } - return nil + return ctx.Err() }); err != nil { - slog.Error("FilteredRelationshipSink Error", attr.Error(err)) + s.recordError(fmt.Errorf("%s: inserting relationships: %w", s.operationName, err)) } } +// recordError cancels the pipeline so submissions cannot block after the writer stops. +func (s *FilteredRelationshipSink) recordError(err error) { + if contextError := s.ctx.Err(); contextError != nil && errors.Is(err, contextError) { + return + } + + s.errorLock.Lock() + if s.workerError == nil { + s.workerError = err + slog.ErrorContext(s.ctx, "FilteredRelationshipSink Error", attr.Error(err)) + } + s.errorLock.Unlock() + s.cancel() +} + // deltaFilterWorker filters out duplicate edges before they reach the insert worker. It checks whether // an edge has already been tracked in the delta subgraph; if not, it forwards it to the insert channel. func (s *FilteredRelationshipSink) deltaFilterWorker(ctx context.Context, filterC, insertC chan EnsureRelationshipJob) { @@ -148,8 +180,7 @@ func (s *FilteredRelationshipSink) deleteMissingEdges(ctx context.Context) error return nil } -// worker is the main goroutine responsible for managing the entire lifecycle of the sink. It -// coordinates between filtering, insertion, and deletion phases and handles shutdown gracefully. +// worker coordinates filtering and insertion. Pruning is reserved for successful completion in Done. func (s *FilteredRelationshipSink) worker(ctx context.Context) { defer s.wg.Done() @@ -189,11 +220,6 @@ func (s *FilteredRelationshipSink) worker(ctx context.Context) { close(insertC) insertWG.Wait() - - // Remove any lingering edges after the operation completes - if err := s.deleteMissingEdges(ctx); err != nil { - slog.Error("Failed deleting missing edges after ingestion sink flush", attr.Error(err)) - } } // Stats returns a pointer to the atomic statistics structure tracking processed relationships. @@ -209,11 +235,74 @@ func (s *FilteredRelationshipSink) start(ctx context.Context) { // Submit submits a new job to be processed by the sink. func (s *FilteredRelationshipSink) Submit(ctx context.Context, nextJob EnsureRelationshipJob) bool { - return channels.Submit(ctx, s.jobC, nextJob) + s.submitLock.RLock() + defer s.submitLock.RUnlock() + + if ctx.Err() != nil || s.ctx.Err() != nil { + return false + } + select { + case <-s.submitDone: + return false + default: + } + + select { + case <-ctx.Done(): + return false + case <-s.ctx.Done(): + return false + case <-s.submitDone: + return false + case s.jobC <- nextJob: + return true + } +} + +// shutdown stops submissions and waits for workers. The first shutdown call determines whether to prune. +func (s *FilteredRelationshipSink) shutdown(prune bool) { + s.shutdownOnce.Do(func() { + defer s.cancel() + if !prune { + s.cancel() + } + + close(s.submitDone) + s.submitLock.Lock() + close(s.jobC) + s.submitLock.Unlock() + s.wg.Wait() + + s.errorLock.Lock() + s.doneError = s.workerError + s.errorLock.Unlock() + + if s.doneError == nil { + if err := s.ctx.Err(); err != nil { + s.doneError = err + } else if prune { + s.doneError = s.deleteMissingEdges(s.ctx) + if s.doneError != nil { + slog.ErrorContext(s.ctx, "Error deleting stale relationships", attr.Error(s.doneError)) + } + } + } + }) +} + +// Done drains submitted jobs and removes obsolete relationships only if processing succeeded. +// Repeated calls return the same result. Producers must complete successfully before calling Done. +func (s *FilteredRelationshipSink) Done() error { + s.shutdown(true) + return s.doneError } -// Done signals the end of processing and waits for all workers to complete. -func (s *FilteredRelationshipSink) Done() { - close(s.jobC) - s.wg.Wait() +// Abort cancels processing and waits for workers without deleting obsolete relationships. +// It returns any writer error, excluding cancellation caused by shutdown. Repeated calls are safe. +// Calling Done after Abort returns a cancellation or worker error. +func (s *FilteredRelationshipSink) Abort() error { + s.shutdown(false) + s.errorLock.Lock() + defer s.errorLock.Unlock() + return s.workerError } diff --git a/packages/go/analysis/post/sink_callers_test.go b/packages/go/analysis/post/sink_callers_test.go new file mode 100644 index 000000000000..038e19394081 --- /dev/null +++ b/packages/go/analysis/post/sink_callers_test.go @@ -0,0 +1,251 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package post_test + +import ( + "context" + "errors" + "testing" + + adAnalysis "github.com/specterops/bloodhound/packages/go/analysis/ad" + azureAnalysis "github.com/specterops/bloodhound/packages/go/analysis/azure" + "github.com/specterops/bloodhound/packages/go/analysis/post" + "github.com/specterops/bloodhound/packages/go/graphschema/ad" + "github.com/specterops/bloodhound/packages/go/graphschema/azure" + "github.com/specterops/dawgs/cypher/models/cypher" + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/util/size" + "github.com/stretchr/testify/require" +) + +func TestFilteredRelationshipSink_CallersCompleteOnlySuccessfulGenerations(t *testing.T) { + var processors = []struct { + name string + trackedKind graph.Kind + producerKind graph.Kind + process func(context.Context, graph.Database) (*post.AtomicPostProcessingStats, error) + }{ + { + name: "owns enumeration", trackedKind: ad.Owns, producerKind: ad.OwnsRaw, + process: func(ctx context.Context, database graph.Database) (*post.AtomicPostProcessingStats, error) { + return adAnalysis.PostOwnsAndWriteOwner(ctx, database, &adAnalysis.LocalGroupData{}) + }, + }, + { + name: "write owner enumeration", trackedKind: ad.WriteOwner, producerKind: ad.WriteOwnerRaw, + process: func(ctx context.Context, database graph.Database) (*post.AtomicPostProcessingStats, error) { + return adAnalysis.PostOwnsAndWriteOwner(ctx, database, &adAnalysis.LocalGroupData{}) + }, + }, + { + name: "app role assignments", trackedKind: azure.AZMGAddSecret, producerKind: azure.Contains, + process: azureAnalysis.AppRoleAssignments, + }, + { + name: "add owner", trackedKind: azure.AddOwner, producerKind: azure.Contains, + process: azureAnalysis.CreateAZAddOwnerEdge, + }, + { + name: "user role assignments", trackedKind: azure.ResetPassword, producerKind: azure.Contains, + process: azureAnalysis.UserRoleAssignments, + }, + } + for _, processor := range processors { + t.Run(processor.name, func(t *testing.T) { + for _, outcome := range []string{"producer failure", "successful empty generation", "writer failure", "pruning failure"} { + t.Run(outcome, func(t *testing.T) { + var ( + failure = errors.New(outcome) + database = &sinkCallerTestDatabase{ + sinkTestDatabase: &sinkTestDatabase{relationships: map[graph.ID]bool{77: true}}, + trackedKind: processor.trackedKind, producerKind: processor.producerKind, + } + ) + switch outcome { + case "producer failure": + database.producerError = failure + case "writer failure": + database.flushError = failure + case "pruning failure": + database.deleteError = failure + } + + _, err := processor.process(t.Context(), database) + if outcome == "successful empty generation" { + require.NoError(t, err) + require.Empty(t, database.relationships, "a complete empty result must remove obsolete edges") + } else { + require.ErrorIs(t, err, failure) + require.True(t, database.relationships[77], "incomplete processing must retain the existing edge") + } + if outcome == "producer failure" { + require.True(t, database.failureInjected, "the failure must occur in production after loading the tracker") + } + }) + } + }) + } +} + +func TestFilteredRelationshipSink_UserRoleAssignmentsRetainEdgesOnInvalidGroupProperty(t *testing.T) { + var database = &sinkCallerTestDatabase{ + sinkTestDatabase: &sinkTestDatabase{relationships: map[graph.ID]bool{77: true}}, + trackedKind: azure.ResetPassword, + invalidGroup: true, + } + _, err := azureAnalysis.UserRoleAssignments(t.Context(), database) + require.Error(t, err) + require.True(t, database.relationships[77], "an invalid AddMembers input must not prune existing role edges") +} + +// Query kinds select the producer fault so prerequisite and tracker reads still succeed. +type sinkCallerTestDatabase struct { + *sinkTestDatabase + trackedKind graph.Kind + producerKind graph.Kind + producerError error + failureInjected bool + invalidGroup bool +} + +func (s *sinkCallerTestDatabase) ReadTransaction(_ context.Context, delegate graph.TransactionDelegate, _ ...graph.TransactionOption) error { + return delegate(&sinkCallerTransaction{database: s}) +} + +type sinkCallerTransaction struct { + graph.Transaction + database *sinkCallerTestDatabase +} + +func (s *sinkCallerTransaction) Nodes() graph.NodeQuery { + return &sinkCallerNodeQuery{} +} + +func (s *sinkCallerTransaction) Relationships() graph.RelationshipQuery { + return &sinkCallerRelationshipQuery{database: s.database} +} + +func (s *sinkCallerTransaction) GraphQueryMemoryLimit() size.Size { + return 0 +} + +type sinkCallerRelationshipQuery struct { + graph.RelationshipQuery + database *sinkCallerTestDatabase + criteria graph.Criteria +} + +func (s *sinkCallerRelationshipQuery) Filter(criteria graph.Criteria) graph.RelationshipQuery { + s.criteria = criteria + return s +} + +func (s *sinkCallerRelationshipQuery) Filterf(provider graph.CriteriaProvider) graph.RelationshipQuery { + return s.Filter(provider()) +} + +func (s *sinkCallerRelationshipQuery) producerFailure() error { + if s.database.producerError != nil && sinkCallerCriteriaHasKind(s.criteria, s.database.producerKind) { + s.database.failureInjected = true + return s.database.producerError + } + return nil +} + +func (s *sinkCallerRelationshipQuery) Fetch(delegate func(graph.Cursor[*graph.Relationship]) error) error { + if err := s.producerFailure(); err != nil { + return err + } + if sinkCallerCriteriaHasKind(s.criteria, s.database.trackedKind) { + return delegate(newSinkCallerCursor(graph.NewRelationship(77, 1, 2, graph.NewProperties(), s.database.trackedKind))) + } + return delegate(newSinkCallerCursor[*graph.Relationship]()) +} + +func (s *sinkCallerRelationshipQuery) FetchIDs(delegate func(graph.Cursor[graph.ID]) error) error { + return delegate(newSinkCallerCursor[graph.ID]()) +} + +func (s *sinkCallerRelationshipQuery) FetchDirection(_ graph.Direction, delegate func(graph.Cursor[graph.DirectionalResult]) error) error { + if err := s.producerFailure(); err != nil { + return err + } + if s.database.invalidGroup && sinkCallerCriteriaHasKind(s.criteria, azure.Contains) && sinkCallerCriteriaHasKind(s.criteria, azure.Group) { + var properties = graph.NewProperties() + properties.Set(azure.IsAssignableToRole.String(), 42) + return delegate(newSinkCallerCursor(graph.DirectionalResult{ + Node: graph.NewNode(5, properties, azure.Group), + })) + } + return delegate(newSinkCallerCursor[graph.DirectionalResult]()) +} + +type sinkCallerNodeQuery struct { + graph.NodeQuery + criteria graph.Criteria +} + +func (s *sinkCallerNodeQuery) Filter(criteria graph.Criteria) graph.NodeQuery { + s.criteria = criteria + return s +} + +func (s *sinkCallerNodeQuery) Filterf(provider graph.CriteriaProvider) graph.NodeQuery { + return s.Filter(provider()) +} + +func (s *sinkCallerNodeQuery) Fetch(delegate func(graph.Cursor[*graph.Node]) error, _ ...graph.Criteria) error { + if sinkCallerCriteriaHasKind(s.criteria, azure.Tenant) { + return delegate(newSinkCallerCursor(graph.NewNode(3, graph.NewProperties(), azure.Tenant))) + } + return delegate(newSinkCallerCursor[*graph.Node]()) +} + +func (s *sinkCallerNodeQuery) FetchIDs(delegate func(graph.Cursor[graph.ID]) error) error { + return delegate(newSinkCallerCursor[graph.ID]()) +} + +func sinkCallerCriteriaHasKind(criteria graph.Criteria, kind graph.Kind) bool { + switch typedCriteria := criteria.(type) { + case *cypher.KindMatcher: + return typedCriteria.Kinds.ContainsOneOf(kind) + case *cypher.Conjunction: + for _, expression := range typedCriteria.Expressions { + if sinkCallerCriteriaHasKind(expression, kind) { + return true + } + } + } + return false +} + +type sinkCallerCursor[T any] struct { + values chan T +} + +func newSinkCallerCursor[T any](values ...T) *sinkCallerCursor[T] { + var valueChannel = make(chan T, len(values)) + for _, value := range values { + valueChannel <- value + } + close(valueChannel) + return &sinkCallerCursor[T]{values: valueChannel} +} + +func (s *sinkCallerCursor[T]) Chan() chan T { return s.values } +func (s *sinkCallerCursor[T]) Error() error { return nil } +func (s *sinkCallerCursor[T]) Close() {} diff --git a/packages/go/analysis/post/sink_test.go b/packages/go/analysis/post/sink_test.go new file mode 100644 index 000000000000..cda2ec1d948b --- /dev/null +++ b/packages/go/analysis/post/sink_test.go @@ -0,0 +1,249 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package post_test + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/specterops/bloodhound/packages/go/analysis/post" + "github.com/specterops/bloodhound/packages/go/graphschema/common" + "github.com/specterops/dawgs/graph" + "github.com/stretchr/testify/require" +) + +type sinkTestDatabase struct { + graph.Database + graph.Batch + relationships map[graph.ID]bool + created []post.EnsureRelationshipJob + batchError error + createError error + flushError error + deleteError error + writerStarted chan struct{} + writerRelease chan struct{} + flushStarted chan struct{} + flushRelease chan struct{} +} + +func (s *sinkTestDatabase) BatchOperation(ctx context.Context, delegate graph.BatchDelegate, _ ...graph.BatchOption) error { + if s.batchError != nil { + close(s.writerStarted) + <-s.writerRelease + return s.batchError + } + if err := delegate(s); err != nil { + return err + } + if s.flushStarted != nil { + close(s.flushStarted) + s.flushStarted = nil + <-s.flushRelease + } + return s.flushError +} + +func (s *sinkTestDatabase) CreateRelationshipByIDs(fromID, toID graph.ID, kind graph.Kind, properties *graph.Properties) error { + if s.createError != nil { + close(s.writerStarted) + <-s.writerRelease + return s.createError + } + s.created = append(s.created, post.EnsureRelationshipJob{ + FromID: fromID, ToID: toID, Kind: kind, RelProperties: properties.Map, + }) + return nil +} + +func (s *sinkTestDatabase) DeleteRelationship(relationshipID graph.ID) error { + if s.deleteError != nil { + return s.deleteError + } + delete(s.relationships, relationshipID) + return nil +} + +func newSinkTestDatabase() (*sinkTestDatabase, *post.Tracker) { + var ( + database = &sinkTestDatabase{relationships: map[graph.ID]bool{77: true}} + builder = post.NewTrackerBuilder() + ) + builder.TrackEdge(77, 1, 2, graph.StringKind("TrackedRelationship"), nil) + return database, builder.Build() +} + +func waitForSinkResult[T any](t *testing.T, resultChannel <-chan T) T { + t.Helper() + select { + case result := <-resultChannel: + return result + case <-time.After(5 * time.Second): + t.Fatal("sink operation did not finish") + var zero T + return zero + } +} + +func finishSink(t *testing.T, sink *post.FilteredRelationshipSink) error { + t.Helper() + var resultChannel = make(chan error, 1) + go func() { resultChannel <- sink.Done() }() + return waitForSinkResult(t, resultChannel) +} + +func TestFilteredRelationshipSink_DonePrunesEmptySuccessfulGeneration(t *testing.T) { + var ( + database, tracker = newSinkTestDatabase() + sink = post.NewFilteredRelationshipSink(t.Context(), "test", database, tracker) + ) + require.NoError(t, finishSink(t, sink)) + require.Empty(t, database.relationships) + require.NoError(t, sink.Done()) + require.NoError(t, sink.Abort()) + require.False(t, sink.Submit(t.Context(), post.EnsureRelationshipJob{})) +} + +func TestFilteredRelationshipSink_DonePreservesVisitedEdgesAndInsertsNewEdges(t *testing.T) { + var ( + database, tracker = newSinkTestDatabase() + kind = graph.StringKind("TrackedRelationship") + sink = post.NewFilteredRelationshipSink(t.Context(), "test", database, tracker) + ) + require.True(t, sink.Submit(t.Context(), post.EnsureRelationshipJob{FromID: 1, ToID: 2, Kind: kind})) + require.True(t, sink.Submit(t.Context(), post.EnsureRelationshipJob{ + FromID: 3, ToID: 4, Kind: kind, RelProperties: map[string]any{"isacl": true}, + })) + require.NoError(t, finishSink(t, sink)) + require.True(t, database.relationships[77]) + require.Len(t, database.created, 1) + require.Equal(t, graph.ID(3), database.created[0].FromID) + require.Equal(t, graph.ID(4), database.created[0].ToID) + require.Equal(t, true, database.created[0].RelProperties["isacl"]) + require.IsType(t, time.Time{}, database.created[0].RelProperties[common.FirstSeen.String()]) + require.Equal(t, int32(1), *sink.Stats().RelationshipsCreated[kind]) +} + +func TestFilteredRelationshipSink_AbortPreservesUnvisitedEdges(t *testing.T) { + var ( + database, tracker = newSinkTestDatabase() + sink = post.NewFilteredRelationshipSink(t.Context(), "test", database, tracker) + resultChannel = make(chan error, 1) + ) + go func() { + resultChannel <- sink.Abort() + }() + require.NoError(t, waitForSinkResult(t, resultChannel)) + require.True(t, database.relationships[77]) + require.NoError(t, sink.Abort()) + require.ErrorIs(t, sink.Done(), context.Canceled) + require.False(t, sink.Submit(t.Context(), post.EnsureRelationshipJob{})) +} + +func TestFilteredRelationshipSink_ParentCancellationPreservesUnvisitedEdges(t *testing.T) { + var ( + database, tracker = newSinkTestDatabase() + ctx, cancel = context.WithCancel(t.Context()) + sink = post.NewFilteredRelationshipSink(ctx, "test", database, tracker) + ) + cancel() + require.ErrorIs(t, finishSink(t, sink), context.Canceled) + require.True(t, database.relationships[77]) + require.False(t, sink.Submit(t.Context(), post.EnsureRelationshipJob{})) +} + +func TestFilteredRelationshipSink_WriterFailureUnblocksSubmissionsAndPreservesEdges(t *testing.T) { + for _, failureMode := range []string{"batch", "create"} { + t.Run(failureMode, func(t *testing.T) { + var ( + database, tracker = newSinkTestDatabase() + writerError = errors.New("writer failed") + resultChannel = make(chan bool, 1) + ) + database.writerStarted = make(chan struct{}) + database.writerRelease = make(chan struct{}) + if failureMode == "batch" { + database.batchError = writerError + } else { + database.createError = writerError + } + var sink = post.NewFilteredRelationshipSink(t.Context(), "test", database, tracker) + go func() { + for submitted := 0; submitted < 10000; submitted++ { + if !sink.Submit(context.Background(), post.EnsureRelationshipJob{ + FromID: 3, ToID: 4, Kind: graph.StringKind("NewRelationship"), + }) { + resultChannel <- false + return + } + } + resultChannel <- true + }() + waitForSinkResult(t, database.writerStarted) + close(database.writerRelease) + require.False(t, waitForSinkResult(t, resultChannel)) + require.ErrorIs(t, finishSink(t, sink), writerError) + require.True(t, database.relationships[77]) + require.ErrorIs(t, sink.Done(), writerError) + require.ErrorIs(t, sink.Abort(), writerError) + }) + } +} + +func TestFilteredRelationshipSink_FlushFailurePreservesUnvisitedEdges(t *testing.T) { + var ( + database, tracker = newSinkTestDatabase() + writerError = errors.New("batch flush failed") + ) + database.flushError = writerError + var sink = post.NewFilteredRelationshipSink(t.Context(), "test", database, tracker) + require.True(t, sink.Submit(t.Context(), post.EnsureRelationshipJob{ + FromID: 3, ToID: 4, Kind: graph.StringKind("NewRelationship"), + })) + require.ErrorIs(t, finishSink(t, sink), writerError) + require.True(t, database.relationships[77]) +} + +func TestFilteredRelationshipSink_DoneReturnsPruningError(t *testing.T) { + var ( + database, tracker = newSinkTestDatabase() + deleteError = errors.New("deleting relationship failed") + ) + database.deleteError = deleteError + var sink = post.NewFilteredRelationshipSink(t.Context(), "test", database, tracker) + require.ErrorIs(t, finishSink(t, sink), deleteError) + require.True(t, database.relationships[77]) +} + +func TestFilteredRelationshipSink_SubmitDuringShutdownIsRejected(t *testing.T) { + var ( + database, tracker = newSinkTestDatabase() + flushStarted = make(chan struct{}) + resultChannel = make(chan error, 1) + ) + database.flushStarted = flushStarted + database.flushRelease = make(chan struct{}) + var sink = post.NewFilteredRelationshipSink(t.Context(), "test", database, tracker) + go func() { resultChannel <- sink.Done() }() + waitForSinkResult(t, flushStarted) + require.False(t, sink.Submit(context.Background(), post.EnsureRelationshipJob{})) + close(database.flushRelease) + require.NoError(t, waitForSinkResult(t, resultChannel)) + require.Empty(t, database.relationships) +}