diff --git a/internal/graph/fbwriter/segmented.go b/internal/graph/fbwriter/segmented.go new file mode 100644 index 000000000..45e24dc57 --- /dev/null +++ b/internal/graph/fbwriter/segmented.go @@ -0,0 +1,362 @@ +// Package fbwriter — segmented.go implements the bounded-memory SegmentedWriter +// that produces the multi-segment gen-dir graph format the #5901 read substrate +// already understands. It is the write half of the #5890 streaming-write epic +// (issue #5902, slice (c)); the reader (MultiReader / OpenSegmentsWithRanges) +// landed dark in #5901 and is unchanged here. +// +// # The problem it solves (#5726) +// +// The single-file writer (streaming.go) streams entities into ONE growing +// flatbuffers.Builder and finalizes it in a single tmp+rename. The builder +// grows to the FULL serialized size (~150-250 MB on a 60 k-entity corpus, and +// unbounded on huge ones) and the vendored flatbuffers library hard-panics at +// its 2 GiB cap — the #5726 "graph too large to serialize" cliff (today merely +// recovered into a degraded error). SegmentedWriter bounds the builder: it +// FLUSHES a finished segment file and starts a FRESH builder whenever the +// current builder crosses a byte threshold, so PEAK builder memory stays ~one +// threshold, NOT O(graph). +// +// # Layout produced +// +// A graph that fits under the threshold in ONE builder takes the single-file +// fast path: a plain graph..fb written by the existing flat writer, +// byte-identical to today (no dir, no manifest — decision 1). Only a graph too +// large for one builder gets the gen-dir layout: +// +// graph./ +// seg-0000.fb entity segment 0 (entities [MinKey..MaxKey], sorted) +// seg-0001.fb entity segment 1 (entities (prev.MaxKey..MaxKey]) +// ... +// seg-000N.fb relationship segment 0 +// ... +// manifest.json +// +// # Per-segment sort + key bounds (the #5901 forward-carry invariant) +// +// Entities are sorted by their id string with Go's `<` (byte-wise, unsigned) — +// the SAME ordering FlatBuffers `(key)` uses (memcmp on Entity.id) and the SAME +// ordering the reader's KeyRange.contains uses. Because the writer streams the +// GLOBALLY sorted entity slice and cuts a new segment only at a threshold +// boundary, every entity segment is (a) internally sorted, so its FlatBuffers +// EntitiesByKey binary search works, and (b) a CONTIGUOUS, NON-OVERLAPPING +// key window: segment i's MaxKey < segment i+1's MinKey (ids are unique). Each +// SegmentMeta records [MinKey,MaxKey] as the exact byte-wise min (first) and +// max (last) id in that segment. The reader routes a lookup to the one segment +// whose inclusive [Min,Max] window contains the key and provably never +// false-skips a present entity, because the segments tile the whole sorted id +// space with no gaps that could strand a key. See TestSegmentedRoundTrip. +// +// # Independent entity / relationship cadence (decision 5) +// +// Entity segments and relationship segments are flushed on their OWN thresholds +// into SEPARATE files with the matching SegmentKind; from_id/to_id are stable +// id strings (graph.fbs `(key)`-free shared strings) so a relationship in one +// segment resolves its endpoints against entities in any segment via the +// reader's cross-segment LookupEntityByID — no forward-ref bookkeeping. +package fbwriter + +import ( + "fmt" + "os" + "path/filepath" + "sort" + "strconv" + + flatbuffers "github.com/google/flatbuffers/go" + + "github.com/cajasmota/grafel/internal/graph" +) + +// DefaultSegmentThresholdBytes is the default builder-size flush threshold: a +// segment is finalized and a fresh builder started once the in-progress +// builder crosses this many serialized bytes. 100 MB keeps peak builder memory +// far below the flatbuffers 2 GiB cap (#5726) while keeping the segment count +// low on typical corpora. Tunable via GRAFEL_SEGMENT_BYTES. +const DefaultSegmentThresholdBytes = 100 << 20 + +// segInitBuilderSize is the initial capacity of each per-segment builder. It +// only affects the number of internal doubling reallocs, never the finished +// bytes; 4 MB matches NewStreamingWriter's default. +const segInitBuilderSize = 4 << 20 + +// StreamSegmentsEnabled reports whether the flag-gated segmented producer is +// ON (GRAFEL_STREAM_SEGMENTS truthy). OFF by default: producers then write +// today's single flat graph..fb with zero behaviour change. This slice +// does NOT flip the default — the flip rides the final #5912 serve slice. +func StreamSegmentsEnabled() bool { + return envTruthy(os.Getenv("GRAFEL_STREAM_SEGMENTS")) +} + +// segmentThresholdBytes resolves the flush threshold, honouring the +// GRAFEL_SEGMENT_BYTES override (positive integer, bytes) and falling back to +// DefaultSegmentThresholdBytes for an unset/invalid value. +func segmentThresholdBytes() int { + if v := os.Getenv("GRAFEL_SEGMENT_BYTES"); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + return n + } + } + return DefaultSegmentThresholdBytes +} + +// envTruthy interprets an env-var value as a boolean (mirrors the helper in +// internal/graph/algorithms.go; kept local to avoid an import cycle). +func envTruthy(v string) bool { + switch v { + case "1", "t", "T", "true", "TRUE", "True", "yes", "YES", "Yes", "on", "ON", "On": + return true + } + return false +} + +// WriteGraphGenSegmented is the flag-ON producer path. It sorts doc into the +// canonical byte-wise key order, then: +// +// - if the whole graph fits under the threshold in a single builder, takes +// the single-file fast path (writeGraphGenFlat → plain graph..fb, +// byte-identical to the flag-OFF path); otherwise +// - streams entity segments and relationship segments on independent +// thresholds into a fresh graph./ dir, writes manifest.json, flips the +// `current` pointer at the gen dir, and best-effort GCs stale generations. +// +// Returns the absolute gen path written: the graph..fb file for the +// single-file case, or the graph./ dir for a segment-set (filepath.Dir of +// either is stateDir, so directory-keyed sidecar writers land correctly). +func WriteGraphGenSegmented(stateDir string, doc *graph.Document) (string, error) { + if doc == nil { + return "", fmt.Errorf("fbwriter: nil document") + } + sortDocForSegments(doc) + threshold := segmentThresholdBytes() + + // Single-file fast path (decision 1). The probe is bounded: it early-exits + // the instant the builder crosses the threshold, so even a multi-GiB graph + // never materialises more than ~threshold bytes here. + if graphFitsSingleBuilder(doc, threshold) { + return writeGraphGenFlat(stateDir, doc) + } + return writeSegments(stateDir, doc, threshold) +} + +// sortDocForSegments sorts doc into the canonical emission order in place: +// entities by byte-wise id (the FlatBuffers `(key)` / reader KeyRange order), +// relationships by (from,to,kind). Producers already sort before calling, so +// this is idempotent; it is repeated here to make the per-segment key-bounds +// invariant hold regardless of the caller. +func sortDocForSegments(doc *graph.Document) { + sort.SliceStable(doc.Entities, func(i, j int) bool { + return doc.Entities[i].ID < doc.Entities[j].ID + }) + sort.SliceStable(doc.Relationships, func(i, j int) bool { + a, b := &doc.Relationships[i], &doc.Relationships[j] + if a.FromID != b.FromID { + return a.FromID < b.FromID + } + if a.ToID != b.ToID { + return a.ToID < b.ToID + } + return a.Kind < b.Kind + }) +} + +// graphFitsSingleBuilder reports whether the entire graph serializes to fewer +// than threshold bytes in ONE builder. It builds entity then relationship +// tables into a throwaway builder and returns false the instant the running +// serialized size crosses the threshold — a bounded probe whose peak is +// ~threshold+one-record, never O(graph). The builder is discarded; the caller +// re-marshals via the canonical flat path to guarantee byte-identity. +func graphFitsSingleBuilder(doc *graph.Document, threshold int) bool { + b := flatbuffers.NewBuilder(segInitBuilderSize) + for i := range doc.Entities { + buildEntity(b, &doc.Entities[i]) + if int(b.Offset()) >= threshold { + return false + } + } + for i := range doc.Relationships { + buildRelationship(b, &doc.Relationships[i]) + if int(b.Offset()) >= threshold { + return false + } + } + return true +} + +// writeSegments performs the real bounded segmented write: entity segments +// first (each a self-contained Graph FlatBuffer whose entities vector is the +// contiguous sorted slice for that segment), then relationship segments, then +// manifest + pointer flip + GC. Peak builder memory is bounded by the flush +// threshold: each segment builder is finalized, written, and dropped before +// the next is allocated, so no single builder ever holds the whole graph. +func writeSegments(stateDir string, doc *graph.Document, threshold int) (string, error) { + gen := graph.NextGen(stateDir) + genDirName := graph.GenDirName(gen) + genDir := filepath.Join(stateDir, genDirName) + if err := os.MkdirAll(genDir, 0o755); err != nil { + return "", fmt.Errorf("fbwriter: mkdir gen dir %s: %w", genDir, err) + } + + fullMeta := graphMetaFromDoc(doc) + var ( + metas []graph.SegmentMeta + segIdx int + firstSeg = true // the first written segment carries the graph header + ) + + // metaForNextSegment returns the header metadata for the segment about to be + // finalized: the FULL header (repo, git meta, algo stats, communities) on + // the very first segment written — the reader's MultiReader reads header + // fields from segment 0 — and an empty header on every later segment. + metaForNextSegment := func() GraphMetadata { + if firstSeg { + return fullMeta + } + return GraphMetadata{} + } + + // ── entity segments ────────────────────────────────────────────────── + var ( + eb *flatbuffers.Builder + eOffs []flatbuffers.UOffsetT + minKey string + maxKey string + ) + flushEntities := func() error { + if len(eOffs) == 0 { + return nil + } + sw := &StreamingWriter{b: eb, entityOffsets: eOffs} + buf, err := sw.finalizeSafe(metaForNextSegment()) + if err != nil { + return err + } + name := graph.SegmentFileName(segIdx) + if err := writeSegmentFile(genDir, name, buf); err != nil { + return err + } + metas = append(metas, graph.SegmentMeta{ + File: name, + Kind: graph.SegmentEntities, + EntityCount: len(eOffs), + MinKey: minKey, + MaxKey: maxKey, + }) + segIdx++ + firstSeg = false + eb, eOffs, minKey, maxKey = nil, nil, "", "" // free the builder + return nil + } + for i := range doc.Entities { + e := &doc.Entities[i] + if eb == nil { + eb = flatbuffers.NewBuilder(segInitBuilderSize) + } + if len(eOffs) == 0 { + minKey = e.ID // first (== smallest) id in this segment + } + eOffs = append(eOffs, buildEntity(eb, e)) + maxKey = e.ID // last appended (== largest so far) id in this segment + if int(eb.Offset()) >= threshold { + if err := flushEntities(); err != nil { + return "", err + } + } + } + if err := flushEntities(); err != nil { + return "", err + } + + // ── relationship segments (independent cadence) ────────────────────── + var ( + rb *flatbuffers.Builder + rOffs []flatbuffers.UOffsetT + ) + flushRels := func() error { + if len(rOffs) == 0 { + return nil + } + sw := &StreamingWriter{b: rb, relOffsets: rOffs} + buf, err := sw.finalizeSafe(metaForNextSegment()) + if err != nil { + return err + } + name := graph.SegmentFileName(segIdx) + if err := writeSegmentFile(genDir, name, buf); err != nil { + return err + } + metas = append(metas, graph.SegmentMeta{ + File: name, + Kind: graph.SegmentRelationships, + RelCount: len(rOffs), + }) + segIdx++ + firstSeg = false + rb, rOffs = nil, nil // free the builder + return nil + } + for i := range doc.Relationships { + r := &doc.Relationships[i] + if rb == nil { + rb = flatbuffers.NewBuilder(segInitBuilderSize) + } + rOffs = append(rOffs, buildRelationship(rb, r)) + if int(rb.Offset()) >= threshold { + if err := flushRels(); err != nil { + return "", err + } + } + } + if err := flushRels(); err != nil { + return "", err + } + + // Defensive: an empty graph should have taken the single-file fast path; + // if we somehow produced no segments, fall back rather than write an + // invalid (zero-segment) manifest. + if len(metas) == 0 { + return writeGraphGenFlat(stateDir, doc) + } + + m := &graph.Manifest{FormatVersion: graph.ManifestFormatVersion, Segments: metas} + if err := graph.WriteManifest(genDir, m); err != nil { + return "", fmt.Errorf("fbwriter: write manifest: %w", err) + } + if err := graph.WriteCurrentPointerRaw(stateDir, genDirName); err != nil { + return "", fmt.Errorf("fbwriter: flip pointer to %s: %w", genDirName, err) + } + graph.GCStaleGens(stateDir, gen) // best-effort; never fails the write + return genDir, nil +} + +// writeSegmentFile writes buf as genDir/name via a sibling .tmp + rename. The +// gen dir is freshly allocated (NextGen) and its `current` pointer is not +// flipped until every segment + the manifest are durably in place, so a reader +// never resolves the pointer to a half-written segment set. +func writeSegmentFile(genDir, name string, buf []byte) error { + dst := filepath.Join(genDir, name) + tmp := dst + ".tmp" + if err := os.WriteFile(tmp, buf, 0o644); err != nil { + return fmt.Errorf("fbwriter: write segment tmp %s: %w", name, err) + } + if err := os.Rename(tmp, dst); err != nil { + os.Remove(tmp) + return fmt.Errorf("fbwriter: rename segment %s: %w", name, err) + } + return nil +} + +// graphMetaFromDoc lifts the document-level header fields into a GraphMetadata, +// mirroring streamingMarshal's construction so the header segment carries the +// same repo / git / algo / community metadata a single-file graph would. +func graphMetaFromDoc(doc *graph.Document) GraphMetadata { + return GraphMetadata{ + Repo: doc.Repo, + GeneratedAt: doc.GeneratedAt, + IndexedRef: doc.IndexedRef, + IndexedSHA: doc.IndexedSHA, + IsWorktree: doc.IsWorktree, + CoverageStatus: doc.CoverageStatus, + AlgorithmStats: doc.AlgorithmStats, + Communities: doc.Communities, + } +} diff --git a/internal/graph/fbwriter/segmented_test.go b/internal/graph/fbwriter/segmented_test.go new file mode 100644 index 000000000..2cf2cecca --- /dev/null +++ b/internal/graph/fbwriter/segmented_test.go @@ -0,0 +1,323 @@ +package fbwriter_test + +import ( + "fmt" + "os" + "path/filepath" + "testing" + "time" + + "github.com/cajasmota/grafel/internal/graph" + fb "github.com/cajasmota/grafel/internal/graph/fbgraph" + "github.com/cajasmota/grafel/internal/graph/fbreader" + "github.com/cajasmota/grafel/internal/graph/fbwriter" +) + +// bigDoc builds a graph.Document with n entities (unique, sortable ids) and a +// relationship chain, each entity padded with a property so per-entity +// serialized size is large enough that a small threshold forces many segments. +func bigDoc(repo string, n int) *graph.Document { + doc := &graph.Document{ + Version: 1, + GeneratedAt: time.Date(2026, 6, 1, 0, 0, 0, 0, time.UTC), + Repo: repo, + } + doc.Entities = make([]graph.Entity, 0, n) + for i := 0; i < n; i++ { + e := graph.Entity{ + // zero-padded id so lexicographic (byte-wise) order == numeric order. + ID: fmt.Sprintf("ent-%08d", i), + Name: fmt.Sprintf("sym_%08d", i), + QualifiedName: fmt.Sprintf("pkg/mod.sym_%08d", i), + Kind: "function", + SourceFile: fmt.Sprintf("pkg/file_%04d.go", i%1000), + StartLine: i, + Language: "go", + } + e.PropSet("visibility", "public") + e.PropSet("padding", fmt.Sprintf("payload-%08d-xxxxxxxxxxxxxxxxxxxx", i)) + doc.Entities = append(doc.Entities, e) + } + // A relationship chain ent[i] -> ent[i+1], deliberately inserted in REVERSE + // order so the writer's own sort is exercised. + for i := n - 1; i >= 1; i-- { + doc.Relationships = append(doc.Relationships, graph.Relationship{ + ID: fmt.Sprintf("rel-%08d", i), + FromID: fmt.Sprintf("ent-%08d", i-1), + ToID: fmt.Sprintf("ent-%08d", i), + Kind: "calls", + }) + } + doc.Stats.Entities = len(doc.Entities) + doc.Stats.Relationships = len(doc.Relationships) + return doc +} + +// openSegmentSet resolves dir as a segment-set and opens a routed MultiReader +// over it (key ranges attached from the manifest, exactly as a real reader +// would). The caller must Close the returned reader before the TempDir is +// torn down (Windows-safe). +func openSegmentSet(t *testing.T, dir string) (*fbreader.MultiReader, graph.GraphDescriptor) { + t.Helper() + desc, err := graph.CurrentGraphDescriptor(dir) + if err != nil { + t.Fatalf("CurrentGraphDescriptor: %v", err) + } + if desc.Kind != graph.GraphSegmentSet { + t.Fatalf("descriptor kind = %v, want GraphSegmentSet", desc.Kind) + } + ranges := make([]fbreader.KeyRange, 0, len(desc.Manifest.Segments)) + for _, s := range desc.Manifest.Segments { + ranges = append(ranges, fbreader.KeyRange{ + HasEntities: s.EntityCount > 0, + Min: s.MinKey, + Max: s.MaxKey, + }) + } + mr, err := fbreader.OpenSegmentsWithRanges(desc.Segments, ranges) + if err != nil { + t.Fatalf("OpenSegmentsWithRanges: %v", err) + } + return mr, desc +} + +// TestSegmentedRoundTrip — the make-or-break correctness test. A graph large +// enough (with a tiny threshold) to force MULTIPLE segments is written via the +// SegmentedWriter, then read back through the segment-aware, key-routed reader. +// EVERY entity must be findable by id and EVERY relationship present — proving +// the writer's per-segment MinKey/MaxKey bounds never cause the reader to +// false-skip a present entity. +func TestSegmentedRoundTrip(t *testing.T) { + t.Setenv("GRAFEL_STREAM_SEGMENTS", "1") + t.Setenv("GRAFEL_SEGMENT_BYTES", "65536") // 64 KB — forces many segments + dir := t.TempDir() + const n = 3000 + + doc := bigDoc("seg-rt", n) + genPath, err := fbwriter.WriteGraphGen(dir, doc) + if err != nil { + t.Fatalf("WriteGraphGen: %v", err) + } + // A segment set is a gen DIR, not a .fb file. + if filepath.Ext(genPath) == ".fb" { + t.Fatalf("expected a gen dir, got file %q", genPath) + } + + mr, desc := openSegmentSet(t, dir) + defer mr.Close() + + if mr.SegmentCount() < 2 { + t.Fatalf("SegmentCount = %d, want >= 2 (threshold should have split)", mr.SegmentCount()) + } + if mr.EntityCount() != n { + t.Fatalf("EntityCount = %d, want %d", mr.EntityCount(), n) + } + if got := mr.RelationshipCount(); got != n-1 { + t.Fatalf("RelationshipCount = %d, want %d", got, n-1) + } + + // Every entity is findable by id through the ROUTED reader (routing that + // false-skipped a segment would return nil here). + for i := 0; i < n; i++ { + id := fmt.Sprintf("ent-%08d", i) + e := mr.LookupEntityByID(id) + if e == nil { + t.Fatalf("LookupEntityByID(%s) = nil — routing false-skipped a present entity", id) + } + if string(e.Id()) != id { + t.Fatalf("LookupEntityByID(%s) returned id %q", id, e.Id()) + } + } + // A negative lookup: an id outside every segment's range is not found and + // does not spuriously resolve. + if e := mr.LookupEntityByID("ent-99999999"); e != nil { + t.Fatalf("lookup of absent id resolved to %q", e.Id()) + } + + // Every relationship is present. + seen := map[string]bool{} + mr.IterateRelationships(func(r *fb.Relationship) bool { + seen[string(r.FromId())+"->"+string(r.ToId())+":"+string(r.Kind())] = true + return true + }) + for i := 1; i < n; i++ { + key := fmt.Sprintf("ent-%08d->ent-%08d:calls", i-1, i) + if !seen[key] { + t.Fatalf("relationship %q missing after round-trip", key) + } + } + + // Cross-check: the writer's recorded [MinKey,MaxKey] windows must TILE the + // whole id space with no gap and no overlap — the property that guarantees + // the reader never false-skips. Entity segments are contiguous and ordered. + var prevMax string + for _, s := range desc.Manifest.EntitySegments() { + if s.MinKey == "" || s.MaxKey == "" { + t.Fatalf("entity segment %s missing key bounds", s.File) + } + if s.MinKey > s.MaxKey { + t.Fatalf("segment %s has MinKey %q > MaxKey %q", s.File, s.MinKey, s.MaxKey) + } + if prevMax != "" && !(prevMax < s.MinKey) { + t.Fatalf("segment %s MinKey %q overlaps previous MaxKey %q", s.File, s.MinKey, prevMax) + } + prevMax = s.MaxKey + } +} + +// TestSegmentedBoundedBuilder — with a small threshold, N entities produce +// multiple segments each ~<= threshold, and no single segment file holds the +// whole graph. This is the bounded-memory proxy: peak builder ~ one segment. +func TestSegmentedBoundedBuilder(t *testing.T) { + const threshold = 64 * 1024 + t.Setenv("GRAFEL_STREAM_SEGMENTS", "1") + t.Setenv("GRAFEL_SEGMENT_BYTES", fmt.Sprintf("%d", threshold)) + dir := t.TempDir() + const n = 4000 + + if _, err := fbwriter.WriteGraphGen(dir, bigDoc("seg-bound", n)); err != nil { + t.Fatalf("WriteGraphGen: %v", err) + } + desc, err := graph.CurrentGraphDescriptor(dir) + if err != nil { + t.Fatalf("descriptor: %v", err) + } + if desc.Kind != graph.GraphSegmentSet { + t.Fatalf("kind = %v, want segment set", desc.Kind) + } + if len(desc.Segments) < 3 { + t.Fatalf("segment count = %d, want several (>=3)", len(desc.Segments)) + } + // Every segment file is bounded near the threshold. Each segment finalizes + // only AFTER a record crosses the threshold, so allow the header/vector + // finalize overhead plus one oversized record of slack. + const slack = 64 * 1024 + for _, p := range desc.Segments { + fi, err := os.Stat(p) + if err != nil { + t.Fatalf("stat %s: %v", p, err) + } + if fi.Size() > threshold+slack { + t.Fatalf("segment %s = %d bytes, exceeds threshold %d + slack %d — builder not bounded", + filepath.Base(p), fi.Size(), threshold, slack) + } + } +} + +// TestSingleFileFastPath — a small graph under the threshold takes the flat +// path even with the flag ON: a plain graph..fb, no dir, no manifest, and +// byte-identical to the flag-OFF write of the same doc. +func TestSingleFileFastPath(t *testing.T) { + // Flag-OFF reference write. + offDir := t.TempDir() + offPath, err := fbwriter.WriteGraphGen(offDir, smallDoc("fast")) + if err != nil { + t.Fatalf("flag-off WriteGraphGen: %v", err) + } + offBytes, err := os.ReadFile(offPath) + if err != nil { + t.Fatalf("read off bytes: %v", err) + } + + // Flag-ON write of an identical doc (default large threshold). + t.Setenv("GRAFEL_STREAM_SEGMENTS", "1") + onDir := t.TempDir() + onPath, err := fbwriter.WriteGraphGen(onDir, smallDoc("fast")) + if err != nil { + t.Fatalf("flag-on WriteGraphGen: %v", err) + } + if filepath.Base(onPath) != "graph.1.fb" { + t.Fatalf("flag-on small graph wrote %q, want graph.1.fb (single-file fast path)", filepath.Base(onPath)) + } + // No gen dir, no manifest. + if _, err := os.Stat(filepath.Join(onDir, "graph.1")); err == nil { + t.Fatalf("single-file fast path wrote a gen dir graph.1/") + } + desc, err := graph.CurrentGraphDescriptor(onDir) + if err != nil { + t.Fatalf("descriptor: %v", err) + } + if desc.Kind != graph.GraphSingleFile { + t.Fatalf("kind = %v, want GraphSingleFile", desc.Kind) + } + onBytes, err := os.ReadFile(onPath) + if err != nil { + t.Fatalf("read on bytes: %v", err) + } + if string(offBytes) != string(onBytes) { + t.Fatalf("single-file fast path bytes differ from flag-off write (%d vs %d bytes)", + len(offBytes), len(onBytes)) + } +} + +// TestFlagOffParity — with GRAFEL_STREAM_SEGMENTS unset, even a graph that +// WOULD split under a tiny threshold is written as a single flat file (the +// flag gates the whole segmented path; threshold is irrelevant when OFF). +func TestFlagOffParity(t *testing.T) { + // Threshold set tiny, but the stream flag is UNSET — segmented path must + // not engage. + t.Setenv("GRAFEL_SEGMENT_BYTES", "1024") + os.Unsetenv("GRAFEL_STREAM_SEGMENTS") + dir := t.TempDir() + genPath, err := fbwriter.WriteGraphGen(dir, bigDoc("parity", 2000)) + if err != nil { + t.Fatalf("WriteGraphGen: %v", err) + } + if filepath.Base(genPath) != "graph.1.fb" { + t.Fatalf("flag-off wrote %q, want single flat graph.1.fb", filepath.Base(genPath)) + } + desc, err := graph.CurrentGraphDescriptor(dir) + if err != nil { + t.Fatalf("descriptor: %v", err) + } + if desc.Kind != graph.GraphSingleFile { + t.Fatalf("kind = %v, want GraphSingleFile (flag OFF)", desc.Kind) + } +} + +// TestSegmentedIndependentStreams — entity segments and relationship segments +// are SEPARATE files with the correct Kind and counts, and the per-stream +// totals recorded in the manifest sum to the full graph. +func TestSegmentedIndependentStreams(t *testing.T) { + t.Setenv("GRAFEL_STREAM_SEGMENTS", "1") + t.Setenv("GRAFEL_SEGMENT_BYTES", "65536") + dir := t.TempDir() + const n = 3000 + + if _, err := fbwriter.WriteGraphGen(dir, bigDoc("seg-indep", n)); err != nil { + t.Fatalf("WriteGraphGen: %v", err) + } + desc, err := graph.CurrentGraphDescriptor(dir) + if err != nil { + t.Fatalf("descriptor: %v", err) + } + m := desc.Manifest + entSegs := m.EntitySegments() + relSegs := m.RelationshipSegments() + if len(entSegs) == 0 || len(relSegs) == 0 { + t.Fatalf("want both entity (%d) and relationship (%d) segments", len(entSegs), len(relSegs)) + } + // Streams are disjoint files: no segment carries both entities and rels. + for _, s := range m.Segments { + if s.EntityCount > 0 && s.RelCount > 0 { + t.Fatalf("segment %s mixes streams (ent=%d rel=%d)", s.File, s.EntityCount, s.RelCount) + } + } + // Kind matches the populated count. + for _, s := range entSegs { + if s.Kind != graph.SegmentEntities { + t.Fatalf("entity segment %s has kind %q", s.File, s.Kind) + } + } + for _, s := range relSegs { + if s.Kind != graph.SegmentRelationships { + t.Fatalf("relationship segment %s has kind %q", s.File, s.Kind) + } + } + if m.TotalEntityCount() != n { + t.Fatalf("manifest total entities = %d, want %d", m.TotalEntityCount(), n) + } + if m.TotalRelationshipCount() != n-1 { + t.Fatalf("manifest total rels = %d, want %d", m.TotalRelationshipCount(), n-1) + } +} diff --git a/internal/graph/fbwriter/writer.go b/internal/graph/fbwriter/writer.go index 648b64a95..e00f5771d 100644 --- a/internal/graph/fbwriter/writer.go +++ b/internal/graph/fbwriter/writer.go @@ -73,6 +73,23 @@ func WriteAtomic(outPath string, doc *graph.Document) error { // to the directory-keyed sidecar writer (graph.WriteSidecar keys on // filepath.Dir), so graph-stats.json still lands beside the graph. func WriteGraphGen(stateDir string, doc *graph.Document) (genPath string, err error) { + // #5902 flag-gated producer. When GRAFEL_STREAM_SEGMENTS is ON, route + // through the bounded SegmentedWriter (a single flat file when the graph + // fits under the threshold, a graph./ segment set otherwise). OFF by + // default: fall through to today's single flat-file path, zero behaviour + // change. Gating here wires every producer call site (full-index, + // incremental, links, enrichment write-back) through one switch. + if StreamSegmentsEnabled() { + return WriteGraphGenSegmented(stateDir, doc) + } + return writeGraphGenFlat(stateDir, doc) +} + +// writeGraphGenFlat is the single-flat-file producer core: marshal the whole +// doc into one buffer (fail-softing an oversized-graph panic into an error) and +// write it as a new graph..fb generation. It is the flag-OFF path and the +// SegmentedWriter's single-file fast path, so both share byte-identical output. +func writeGraphGenFlat(stateDir string, doc *graph.Document) (genPath string, err error) { buf, err := Marshal(doc) if err != nil { return "", fmt.Errorf("fbwriter: marshal: %w", err) diff --git a/internal/graph/genpath.go b/internal/graph/genpath.go index ebc240d6d..ea7a2f7a4 100644 --- a/internal/graph/genpath.go +++ b/internal/graph/genpath.go @@ -282,16 +282,30 @@ func parseSegPointer(raw string) (genDirName string, ok bool) { // sequence monotonic even if the pointer was manually deleted while gen files // remain, so a stale reader never resolves a freshly-written gen to an older // file. +// #5902: the segmented writer emits a generation as a graph./ DIR (not a +// graph..fb file), and points `current` at that dir. NextGen therefore +// scans the UNION of gen files and gen dirs — and reads the pointer via the +// segment-aware readPointerRaw so a segment-set `current` value is counted too +// — keeping the sequence monotonic across a mixed single-file/segment-set +// history. Trusting only the file-shaped pointer/entries (as before) would let +// a fresh segmented write collide with the immediately-previous gen dir. func NextGen(dir string) uint64 { var maxGen uint64 - if name, ok := readPointer(dir); ok { - if v, _ := parseGen(name); v > maxGen { + if raw, ok := readPointerRaw(dir); ok { + if v, single := parseGen(raw); single && v > maxGen { maxGen = v + } else if name, seg := parseSegPointer(raw); seg { + if v, ok := parseGenDir(name); ok && v > maxGen { + maxGen = v + } } } if ents, err := os.ReadDir(dir); err == nil { for _, e := range ents { if e.IsDir() { + if v, ok := parseGenDir(e.Name()); ok && v > maxGen { + maxGen = v + } continue } if v, ok := parseGen(e.Name()); ok && v > maxGen {