From c3207d596379b8e5b33949f0dbfa76dead67aae0 Mon Sep 17 00:00:00 2001 From: Ted Xu Date: Wed, 27 May 2026 16:42:27 +0800 Subject: [PATCH] fix: make BulkPackWriterV3.Write atomic with single manifest commit (#49711) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit issue: #49710 related: #49069 BulkPackWriterV3.Write currently bumps the loon manifest version up to four times per call (inserts, stats, delta, bm25 each open and commit their own transaction). That breaks atomicity for concurrent readers, leaves orphan manifest versions behind on partial failure, and churns loon's manifest history. This PR refactors the V3 sync write path into a do-then-commit shape: Phase 1 (slow, runs once outside the retry loop): writeInserts, writeStats, writeDelta, writeBM25Stasts write all the data files (parquet/LOB via the loon writer, deltalog, stat blobs) and each helper returns its contribution (ColumnGroups, SegmentOutput, StatEntry list, DeltaLogEntry list) without touching shared state. Phase 2 (fast, retried on transient loon errors): Write assembles one packed.ManifestUpdates from the helper returns and calls a single new primitive packed.CommitManifestUpdates, which opens a short loon transaction, stages every change, and commits once. To make this clean the loon FFI writers are decoupled from manifest concerns entirely: FFIPackedWriter and FFISegmentWriter no longer take a baseVersion at construction. Their Close primitives return data carriers (packed.ColumnGroups, packed.SegmentOutput) that own the C output. packed.CommitManifestUpdates is the only Go-level entry point that mutates a manifest. Empty payload short-circuits to the unchanged manifest path without opening a transaction. The same do-then-commit shape applies to the V2-style binlog import writers (PackedManifestRecordWriter, PackedTextManifestRecordWriter): they assemble inserts plus bloom-filter and BM25 stat blobs into one ManifestUpdates and call CommitManifestUpdates once. Internal renames: the low-level adapters packedRecordManifestWriter and packedTextManifestWriter are renamed to packedRecordBatchWriter and packedTextBatchWriter because they don't touch the manifest themselves — they translate storage.Record into arrow.Record and track per-column-group sizes. The interface manifestRecordWriter is renamed to packedBatchWriter to match. New tests in pack_writer_v3_test.go: TestWrite_SingleVersionBumpAcrossSections drives an insert + delta + flush so all four sections fire and asserts the new manifest version equals baseVer + 1. TestWrite_RetryDoesNotLeakVersionBumps mockey-patches CommitManifestUpdates to fail once with packed.ErrLoonTransient and asserts the eventual successful commit still yields baseVer + 1 rather than baseVer + N_retries. All 11 V3 tests pass. Existing packed and storage tests pass. Out of scope (filed as follow-ups in #49710's discussion): backfill_compactor.go's runBm25Function still issues a separate AddStatsToManifest after the column-groups commit, producing two version bumps per backfill. The architecture supports folding the two but the change isn't strictly needed for the V3 sync atomicity fix. FFIPackedWriter lacks a Destroy method; on loon_writer_close failure the handle and cProperties leak. Pre-existing on master. --------- Signed-off-by: Ted Xu --- internal/compaction/common_test.go | 9 +- .../datanode/compactor/backfill_compactor.go | 35 +- .../flushcommon/syncmgr/pack_writer_v3.go | 404 ++++++++---------- .../syncmgr/pack_writer_v3_helpers_test.go | 87 ++++ .../syncmgr/pack_writer_v3_test.go | 110 +++++ internal/storage/binlog_record_writer.go | 155 +++---- .../storage/binlog_record_writer_v3_test.go | 77 ++++ internal/storage/record_writer.go | 120 +++--- internal/storagev2/packed/manifest_commit.go | 158 +++++++ .../storagev2/packed/manifest_commit_test.go | 295 +++++++++++++ .../packed/packed_reader_ffi_test.go | 65 +-- .../storagev2/packed/packed_writer_ffi.go | 124 ++++-- .../packed/packed_writer_ffi_test.go | 65 ++- .../storagev2/packed/segment_writer_ffi.go | 165 +++---- internal/storagev2/packed/transaction_test.go | 9 +- internal/storagev2/packed/type.go | 13 +- 16 files changed, 1357 insertions(+), 534 deletions(-) create mode 100644 internal/flushcommon/syncmgr/pack_writer_v3_helpers_test.go create mode 100644 internal/storage/binlog_record_writer_v3_test.go create mode 100644 internal/storagev2/packed/manifest_commit.go create mode 100644 internal/storagev2/packed/manifest_commit_test.go diff --git a/internal/compaction/common_test.go b/internal/compaction/common_test.go index 3dffd92d42..4d43d2281c 100644 --- a/internal/compaction/common_test.go +++ b/internal/compaction/common_test.go @@ -386,7 +386,7 @@ func (s *CommonSuite) createBaseManifest(basePath string, storageConfig *indexpb {Columns: []int{0, 1}, GroupID: storagecommon.DefaultShortColumnGroupID}, } - pw, err := packed.NewFFIPackedWriter(basePath, 0, arrowSchema, columnGroups, storageConfig, nil) + pw, err := packed.NewFFIPackedWriter(basePath, arrowSchema, columnGroups, storageConfig, nil) require.NoError(t, err) // Write minimal data to create a valid manifest @@ -401,7 +401,12 @@ func (s *CommonSuite) createBaseManifest(basePath string, storageConfig *indexpb err = pw.WriteRecordBatch(rec) require.NoError(t, err) - manifestPath, err := pw.Close() + out, err := pw.Close() + require.NoError(t, err) + defer out.Destroy() + + manifestPath, err := packed.CommitManifestUpdates(basePath, packed.ManifestEarliest, storageConfig, + &packed.ManifestUpdates{NewFiles: out}) require.NoError(t, err) return manifestPath diff --git a/internal/datanode/compactor/backfill_compactor.go b/internal/datanode/compactor/backfill_compactor.go index b40606452b..dd39f0f851 100644 --- a/internal/datanode/compactor/backfill_compactor.go +++ b/internal/datanode/compactor/backfill_compactor.go @@ -63,11 +63,16 @@ type backfillWriter interface { Manifest() string } -// ffiWriterWrapper wraps packed.FFIPackedWriter to implement backfillWriter, -// capturing the manifest string returned by Close. +// ffiWriterWrapper wraps packed.FFIPackedWriter to implement backfillWriter. +// FFIPackedWriter itself knows nothing about manifests; this wrapper closes +// the writer to collect column groups, then commits them to the existing +// manifest via packed.CommitManifestUpdates. type ffiWriterWrapper struct { - writer *packed.FFIPackedWriter - manifest string + writer *packed.FFIPackedWriter + basePath string + baseVersion int64 + storageConfig *indexpb.StorageConfig + manifest string } func (w *ffiWriterWrapper) WriteRecordBatch(r arrow.Record) error { @@ -75,11 +80,20 @@ func (w *ffiWriterWrapper) WriteRecordBatch(r arrow.Record) error { } func (w *ffiWriterWrapper) Close() error { - manifest, err := w.writer.Close() + out, err := w.writer.Close() if err != nil { return err } - w.manifest = manifest + if out == nil { + return nil + } + defer out.Destroy() + newPath, err := packed.CommitManifestUpdates(w.basePath, w.baseVersion, w.storageConfig, + &packed.ManifestUpdates{NewFiles: out}) + if err != nil { + return err + } + w.manifest = newPath return nil } @@ -389,7 +403,7 @@ func (t *backfillCompactionTask) setupWriter(outputField *schemapb.FieldSchema, if err != nil { return nil, merr.WrapErrServiceInternal("failed to parse existing manifest for V3 backfill", err.Error()) } - ffiWriter, err := packed.NewFFIPackedWriter(basePath, existingVersion, arrowSchema, newColumnGroups, t.compactionParams.StorageConfig, pluginContext) + ffiWriter, err := packed.NewFFIPackedWriter(basePath, arrowSchema, newColumnGroups, t.compactionParams.StorageConfig, pluginContext) if err != nil { return nil, err } @@ -397,7 +411,12 @@ func (t *backfillCompactionTask) setupWriter(outputField *schemapb.FieldSchema, // Use AddColumnGroup semantics so Loon does not require the count to match // the existing groups in the manifest. ffiWriter.AsNewColumnGroups() - result.writer = &ffiWriterWrapper{writer: ffiWriter} + result.writer = &ffiWriterWrapper{ + writer: ffiWriter, + basePath: basePath, + baseVersion: existingVersion, + storageConfig: t.compactionParams.StorageConfig, + } result.basePath = basePath } else { // V2: use PackedWriter with explicit file paths. diff --git a/internal/flushcommon/syncmgr/pack_writer_v3.go b/internal/flushcommon/syncmgr/pack_writer_v3.go index fdf13d233b..ffc9f24c33 100644 --- a/internal/flushcommon/syncmgr/pack_writer_v3.go +++ b/internal/flushcommon/syncmgr/pack_writer_v3.go @@ -33,7 +33,6 @@ import ( "github.com/milvus-io/milvus/internal/storagev2/packed" "github.com/milvus-io/milvus/pkg/v3/log" "github.com/milvus-io/milvus/pkg/v3/proto/datapb" - "github.com/milvus-io/milvus/pkg/v3/proto/indexcgopb" "github.com/milvus-io/milvus/pkg/v3/proto/indexpb" "github.com/milvus-io/milvus/pkg/v3/util/metautil" "github.com/milvus-io/milvus/pkg/v3/util/paramtable" @@ -41,14 +40,17 @@ import ( "github.com/milvus-io/milvus/pkg/v3/util/typeutil" ) -// manifestRecordWriter is the common interface for both packedRecordManifestWriter -// and packedTextManifestWriter used for V3 storage writes. -type manifestRecordWriter interface { - storage.RecordWriter +// packedBatchWriter is the common subset both storage.packedRecordBatchWriter +// and storage.packedTextBatchWriter expose to the V3 sync path. Close +// returns a packed.WriterOutput regardless of which concrete writer is +// in use; the V3 sync path treats both uniformly. +type packedBatchWriter interface { + Write(r storage.Record) error + Close() (packed.WriterOutput, error) + GetWrittenUncompressed() uint64 GetColumnGroupWrittenCompressed(columnGroup typeutil.UniqueID) uint64 GetColumnGroupWrittenUncompressed(columnGroup typeutil.UniqueID) uint64 GetWrittenPaths(columnGroup typeutil.UniqueID) string - GetWrittenManifest() string GetWrittenRowNum() int64 } @@ -57,23 +59,22 @@ type BulkPackWriterV3 struct { manifestPath string - // initialManifestPath captures the manifest path observed when Write was - // first invoked. resetForRetry restores manifestPath to this value so each - // retry attempt restarts from the same base manifest version. Touching this - // invariant breaks the retry correctness argument — see - // ccmd/pack_writer_v3_retry_plan.md. + // initialManifestPath captures the manifest path observed when Write is + // invoked. Read-based merges (existing bloom filter / BM25 file lists) + // reference this stable value instead of bw.manifestPath so they cannot + // drift after Phase 2's commit updates manifestPath. initialManifestPath string // pendingMetaCacheActions accumulates RollStats / MergeBm25Stats actions // that writeStats / writeBM25Stasts would otherwise apply directly to the - // metaCache. They are deferred and drained by SyncTask.Run only after a - // successful Write, so that retried attempts do not double-apply. + // metaCache. They are drained by SyncTask.Run only after a successful + // Write, so a failed commit retry never observes partial metaCache state. pendingMetaCacheActions []metacache.SegmentAction - // singlePKStats holds the per-batch PK statistic computed once per Write - // call. It is reused across retry attempts and passed explicitly into the - // serializer's *With variants so merged-stats computation can include this - // batch without requiring the metaCache to be updated yet. + // singlePKStats holds the per-batch PK statistic computed during + // writeStats and passed explicitly into the serializer's *With variants + // so merged-stats computation can include this batch without requiring + // the metaCache to be updated yet. singlePKStats *storage.PrimaryKeyStats } @@ -89,17 +90,20 @@ func NewBulkPackWriterV3(metaCache metacache.MetaCache, schema *schemapb.Collect } } -// Write performs the four manifest-mutating steps for a SyncPack inside a -// retry loop. The retry handles loon transaction conflicts (currently -// surfaced as packed.ErrLoonTransient — see internal/storagev2/packed -// /ffi_common.go for the reason this is coarse-grained today). +// Write executes a SyncPack as do-then-commit: // -// Each attempt restarts from bw.initialManifestPath via resetForRetry so the -// manifest read base does not drift across retries. metaCache mutations -// produced by writeStats / writeBM25Stasts are accumulated in -// bw.pendingMetaCacheActions and applied via a deferred drainer that runs -// only when Write returns nil — this is what makes retries idempotent on -// metaCache without exposing the pending list to callers. +// 1. Phase 1 (slow, runs once): write all data files — parquet/LOB via the +// loon writer, deltalog via the deltalog writer, stat blobs via the +// filesystem FFI — and assemble a single packed.ManifestUpdates payload +// describing every change the segment needs to register. +// 2. Phase 2 (fast, retried on transient loon errors): call +// packed.CommitManifestUpdates once. The loon transaction handle is +// opened, all changes are staged, and the transaction is committed in +// one shot, producing exactly one manifest version bump. +// +// metaCache mutations produced by Phase 1's stats helpers are accumulated +// in bw.pendingMetaCacheActions and only drained on success, so a failed +// retry cycle leaves the metaCache untouched. func (bw *BulkPackWriterV3) Write(ctx context.Context, pack *SyncPack) ( inserts map[int64]*datapb.FieldBinlog, deltas *datapb.FieldBinlog, @@ -112,9 +116,9 @@ func (bw *BulkPackWriterV3) Write(ctx context.Context, pack *SyncPack) ( log := log.Ctx(ctx) bw.initialManifestPath = bw.manifestPath - // Drain deferred metaCache actions on successful Write only. If the retry - // loop terminates with an error we leave the metaCache untouched, so a - // failed sync does not pollute bloom-filter / BM25 history. + // Drain deferred metaCache actions on successful Write only. If we + // return an error we leave the metaCache untouched, so a failed sync + // does not pollute bloom-filter / BM25 history. defer func() { if err != nil { return @@ -128,41 +132,66 @@ func (bw *BulkPackWriterV3) Write(ctx context.Context, pack *SyncPack) ( ) }() - var deltaSummary *datapb.FieldBinlog + basePath, baseVersion, parseErr := packed.UnmarshalManifestPath(bw.initialManifestPath) + if parseErr != nil { + err = parseErr + return + } + + // Phase 1: write files. Each helper returns its contribution to the + // final ManifestUpdates instead of mutating shared state. + var ( + insertFiles packed.WriterOutput + statEntries []packed.StatEntry + deltaEntries []packed.DeltaLogEntry + bm25Entries []packed.StatEntry + ) + defer func() { + if insertFiles != nil { + insertFiles.Destroy() + } + }() + + if inserts, insertFiles, err = bw.writeInserts(ctx, pack, basePath); err != nil { + log.Warn("failed to write insert data", zap.Error(err)) + return + } + if stats, statEntries, err = bw.writeStats(ctx, pack, basePath); err != nil { + log.Warn("failed to process stats blob", zap.Error(err)) + return + } + if deltas, deltaEntries, err = bw.writeDelta(ctx, pack, basePath); err != nil { + log.Warn("failed to process delta blob", zap.Error(err)) + return + } + if bm25Stats, bm25Entries, err = bw.writeBM25Stasts(ctx, pack, basePath); err != nil { + log.Warn("failed to process bm25 stats blob", zap.Error(err)) + return + } + + updates := &packed.ManifestUpdates{ + NewFiles: insertFiles, + DeltaLogs: deltaEntries, + Stats: append(statEntries, bm25Entries...), + } + + // Phase 2: commit the assembled updates. CommitManifestUpdates short- + // circuits to the unchanged manifest path when updates carry nothing. + // The outer retry.Do handles transient FFI errors classified as + // packed.ErrLoonTransient; loon's own optimistic retry covers + // manifest-version conflicts within a single attempt. err = retry.Do(ctx, func() error { - bw.resetForRetry() - - var innerErr error - if inserts, manifest, innerErr = bw.writeInserts(ctx, pack); innerErr != nil { - log.Warn("failed to write insert data", zap.Error(innerErr)) - return classifyLoonErr(innerErr) - } - // Update manifestPath after writeInserts - bw.manifestPath = manifest - - // writeStats for V3 adds bloom filter stats to manifest - if stats, innerErr = bw.writeStats(ctx, pack); innerErr != nil { - log.Warn("failed to process stats blob", zap.Error(innerErr)) - return classifyLoonErr(innerErr) - } - // writeDelta for V3 updates manifest and returns delta summary - if manifest, deltaSummary, innerErr = bw.writeDelta(ctx, pack); innerErr != nil { - log.Warn("failed to process delta blob", zap.Error(innerErr)) - return classifyLoonErr(innerErr) - } - // writeBM25Stasts for V3 adds BM25 stats to manifest - if bm25Stats, innerErr = bw.writeBM25Stasts(ctx, pack); innerErr != nil { - log.Warn("failed to process bm25 stats blob", zap.Error(innerErr)) - return classifyLoonErr(innerErr) + newPath, commitErr := packed.CommitManifestUpdates(basePath, baseVersion, bw.storageConfig, updates) + if commitErr != nil { + return classifyLoonErr(commitErr) } + bw.manifestPath = newPath return nil }, bw.writeRetryOpts...) if err != nil { return } - // For V3, stats are in manifest; delta summary is returned for compaction trigger - deltas = deltaSummary manifest = bw.manifestPath size = bw.sizeWritten return @@ -186,70 +215,69 @@ func classifyLoonErr(err error) error { return retry.Unrecoverable(err) } -// resetForRetry restores per-attempt state. It MUST restore manifestPath to -// initialManifestPath — see the field doc on initialManifestPath for why. -func (bw *BulkPackWriterV3) resetForRetry() { - bw.manifestPath = bw.initialManifestPath - bw.sizeWritten = 0 - bw.pendingMetaCacheActions = bw.pendingMetaCacheActions[:0] - bw.singlePKStats = nil - // Files written under _stats/ and _delta/ by previous failed attempts are - // intentionally NOT cleaned up here: they are unreferenced by any - // committed manifest and will be reclaimed by loon GC. See - // ccmd/pack_writer_v3_retry_plan.md decision (2). -} - -func (bw *BulkPackWriterV3) writeInserts(ctx context.Context, pack *SyncPack) (map[int64]*datapb.FieldBinlog, string, error) { +// writeInserts writes the insert data files for the SyncPack and returns +// the produced WriterOutput plus the per-column-group binlog metadata. +// The returned WriterOutput is owned by the caller and must be Destroy'd +// after the surrounding commit (success or failure). +func (bw *BulkPackWriterV3) writeInserts(ctx context.Context, pack *SyncPack, basePath string) (logs map[int64]*datapb.FieldBinlog, files packed.WriterOutput, err error) { if len(pack.insertData) == 0 { - return make(map[int64]*datapb.FieldBinlog), bw.manifestPath, nil + return make(map[int64]*datapb.FieldBinlog), nil, nil } rec, err := bw.serializeBinlog(ctx, pack) if err != nil { - return nil, "", err + return nil, nil, err } defer rec.Release() tsFrom, tsTo := bw.getTsRange(rec) pluginContextPtr := bw.getPluginContext(pack.collectionID) - // NOTE: this used to wrap the call below in its own retry.Do; the outer - // BulkPackWriterV3.Write retry loop now covers it, so we drop the inner - // retry to avoid Attempts^2 amplification. - logs, manifestPath, err := bw.writeInsertsIntoStorage(ctx, pluginContextPtr, rec, tsFrom, tsTo) + // LOB base path is at partition level: {basePath}/.. = {root}/insert_log/{coll}/{part} + partitionBasePath := path.Dir(basePath) + textColumnConfigs := buildTextColumnConfigs(bw.schema, partitionBasePath) + var w packedBatchWriter + if len(textColumnConfigs) > 0 { + log.Ctx(ctx).Info("using TEXT-aware writer for import", + zap.Int("textFieldCount", len(textColumnConfigs)), + zap.String("basePath", basePath)) + w, err = storage.NewPackedTextBatchWriter("", basePath, bw.schema, + bw.bufferSize, bw.multiPartUploadSize, bw.columnGroups, bw.storageConfig, textColumnConfigs) + } else { + w, err = storage.NewPackedRecordBatchWriter(basePath, bw.schema, + bw.bufferSize, bw.multiPartUploadSize, bw.columnGroups, bw.storageConfig, pluginContextPtr) + } if err != nil { - log.Ctx(ctx).Warn("failed to write inserts into storage", + return nil, nil, err + } + + // Ensure the FFI writer's C resources are reclaimed even if Write fails + // before we reach the explicit Close below. Once Close has been called + // (success or failure) the writer's own defer has already released its + // handle and properties, so closeAttempted gates against a redundant + // second Close from this defer. + closeAttempted := false + defer func() { + if !closeAttempted { + if out, closeErr := w.Close(); closeErr == nil && out != nil { + out.Destroy() + } + } + }() + + if err = w.Write(rec); err != nil { + log.Ctx(ctx).Warn("failed to write inserts", zap.Int64("collectionID", pack.collectionID), zap.Int64("segmentID", pack.segmentID), zap.Error(err)) - return nil, "", err + return nil, nil, err } - return logs, manifestPath, nil -} -func (bw *BulkPackWriterV3) writeInsertsIntoStorage(ctx context.Context, - pluginContextPtr *indexcgopb.StoragePluginContext, - rec storage.Record, - tsFrom typeutil.Timestamp, - tsTo typeutil.Timestamp, -) (map[int64]*datapb.FieldBinlog, string, error) { - log := log.Ctx(ctx) - logs := make(map[int64]*datapb.FieldBinlog) - columnGroups := bw.columnGroups - - var err error - doWrite := func(w manifestRecordWriter) error { - if err = w.Write(rec); err != nil { - if closeErr := w.Close(); closeErr != nil { - log.Error("failed to close writer after write failed", zap.Error(closeErr)) - } - return err - } - // close first the get stats & output - return w.Close() + closeAttempted = true + if files, err = w.Close(); err != nil { + return nil, nil, err } - var manifestPath string getFieldNullCounts := func(columnGroup storagecommon.ColumnGroup) map[int64]int64 { result := make(map[int64]int64, len(columnGroup.Fields)) for _, fieldID := range columnGroup.Fields { @@ -260,31 +288,8 @@ func (bw *BulkPackWriterV3) writeInsertsIntoStorage(ctx context.Context, return result } - basePath, version, err := packed.UnmarshalManifestPath(bw.manifestPath) - if err != nil { - return nil, "", err - } - - // LOB base path is at partition level: {basePath}/.. = {root}/insert_log/{coll}/{part} - partitionBasePath := path.Dir(basePath) - textColumnConfigs := buildTextColumnConfigs(bw.schema, partitionBasePath) - var w manifestRecordWriter - if len(textColumnConfigs) > 0 { - log.Info("using TEXT-aware writer for import", - zap.Int("textFieldCount", len(textColumnConfigs)), - zap.String("basePath", basePath)) - w, err = storage.NewPackedTextManifestWriter("", basePath, version, bw.schema, - bw.bufferSize, bw.multiPartUploadSize, columnGroups, bw.storageConfig, textColumnConfigs) - } else { - w, err = storage.NewPackedRecordManifestWriter(basePath, version, bw.schema, bw.bufferSize, bw.multiPartUploadSize, columnGroups, bw.storageConfig, pluginContextPtr) - } - if err != nil { - return nil, "", err - } - if err = doWrite(w); err != nil { - return nil, "", err - } - for _, columnGroup := range columnGroups { + logs = make(map[int64]*datapb.FieldBinlog) + for _, columnGroup := range bw.columnGroups { columnGroupID := columnGroup.GroupID logs[columnGroupID] = &datapb.FieldBinlog{ FieldID: columnGroupID, @@ -302,75 +307,53 @@ func (bw *BulkPackWriterV3) writeInsertsIntoStorage(ctx context.Context, }, } } - manifestPath = w.GetWrittenManifest() - return logs, manifestPath, nil + return logs, files, nil } -// writeDelta writes deltalog to storage and updates the manifest. -// Returns the updated manifest path and a pathless delta summary FieldBinlog -// containing only EntriesNum and MemorySize for compaction trigger decisions. -func (bw *BulkPackWriterV3) writeDelta(ctx context.Context, pack *SyncPack) (string, *datapb.FieldBinlog, error) { +// writeDelta writes the deltalog file and returns the DeltaLogEntry list +// the caller will fold into ManifestUpdates plus a pathless delta summary +// FieldBinlog (EntriesNum + MemorySize) used by compaction-trigger +// decisions. +func (bw *BulkPackWriterV3) writeDelta(ctx context.Context, pack *SyncPack, basePath string) (*datapb.FieldBinlog, []packed.DeltaLogEntry, error) { if pack.deltaData == nil || pack.deltaData.RowCount == 0 { - return bw.manifestPath, nil, nil + return nil, nil, nil } pkField, err := typeutil.GetPrimaryFieldSchema(bw.schema) if err != nil { - return "", nil, fmt.Errorf("primary key field not found: %w", err) + return nil, nil, fmt.Errorf("primary key field not found: %w", err) } - // Allocate log ID for deltalog logID, err := bw.allocator.AllocOne() if err != nil { - return "", nil, err - } - - // Build deltalog path under basePath/_delta/ - basePath, _, err := packed.UnmarshalManifestPath(bw.manifestPath) - if err != nil { - return "", nil, fmt.Errorf("failed to parse manifest path: %w", err) + return nil, nil, err } deltaPath := metautil.BuildDeltaLogPathV3(basePath, logID) - // Create deltalog writer with V2 storage writer, err := storage.NewDeltalogWriter( ctx, pack.collectionID, pack.partitionID, pack.segmentID, logID, pkField.DataType, deltaPath, storage.WithVersion(storage.StorageV2), storage.WithStorageConfig(bw.storageConfig), ) if err != nil { - return "", nil, fmt.Errorf("failed to create deltalog writer: %w", err) + return nil, nil, fmt.Errorf("failed to create deltalog writer: %w", err) } - // Build Arrow record from delete data using existing utility record, _, _, err := storage.BuildDeleteRecord(pack.deltaData.Pks, pack.deltaData.Tss) if err != nil { - return "", nil, fmt.Errorf("failed to build delete record: %w", err) + return nil, nil, fmt.Errorf("failed to build delete record: %w", err) } defer record.Release() - // Write and close if err := writer.Write(record); err != nil { - return "", nil, fmt.Errorf("failed to write delta record: %w", err) + return nil, nil, fmt.Errorf("failed to write delta record: %w", err) } if err := writer.Close(); err != nil { - return "", nil, fmt.Errorf("failed to close delta writer: %w", err) + return nil, nil, fmt.Errorf("failed to close delta writer: %w", err) } - // Update manifest with the new deltalog - newManifest, err := packed.AddDeltaLogsToManifest( - bw.manifestPath, - bw.storageConfig, - []packed.DeltaLogEntry{{Path: deltaPath, NumEntries: pack.deltaData.RowCount}}, - ) - if err != nil { - return "", nil, fmt.Errorf("failed to add deltalog to manifest: %w", err) - } - - bw.manifestPath = newManifest bw.sizeWritten += pack.deltaData.Size() - // Return delta summary for compaction trigger decisions (no path, only stats) summary := &datapb.FieldBinlog{ Binlogs: []*datapb.Binlog{{ LogID: logID, @@ -378,25 +361,24 @@ func (bw *BulkPackWriterV3) writeDelta(ctx context.Context, pack *SyncPack) (str MemorySize: int64(writer.GetWrittenUncompressed()), }}, } - return newManifest, summary, nil + return summary, []packed.DeltaLogEntry{{Path: deltaPath, NumEntries: pack.deltaData.RowCount}}, nil } -// writeStats overrides the base class to write bloom filter stats into the -// manifest instead of separate binlog files. The stat files are written -// under _stats/ relative to the manifest base path, then registered in -// the manifest via a transaction. -func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack) (map[int64]*datapb.FieldBinlog, error) { +// writeStats writes bloom filter stat blobs under basePath/_stats and +// returns the resulting StatEntry list. The caller folds the entries into +// a ManifestUpdates that commits atomically with inserts / delta / bm25. +func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack, basePath string) (map[int64]*datapb.FieldBinlog, []packed.StatEntry, error) { if len(pack.insertData) == 0 { - return make(map[int64]*datapb.FieldBinlog), nil + return make(map[int64]*datapb.FieldBinlog), nil, nil } serializer, err := NewStorageSerializer(bw.metaCache, bw.schema) if err != nil { - return nil, err + return nil, nil, err } singlePKStats, batchStatsBlob, err := serializer.serializeStatslog(pack) if err != nil { - return nil, err + return nil, nil, err } bw.singlePKStats = singlePKStats @@ -409,10 +391,6 @@ func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack) (map metacache.RollStats(singlePKStats)) pkFieldID := serializer.pkField.GetFieldID() - basePath, _, err := packed.UnmarshalManifestPath(bw.manifestPath) - if err != nil { - return nil, err - } var files []string var memorySize int64 @@ -421,7 +399,7 @@ func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack) (map // loon_transaction_update_stat uses replace semantics, so we must // merge previously written files into the new entry. statKey := fmt.Sprintf("bloom_filter.%d", pkFieldID) - existingStats, err := packed.GetManifestStats(bw.manifestPath, bw.storageConfig) + existingStats, err := packed.GetManifestStats(bw.initialManifestPath, bw.storageConfig) if err == nil { if existing, ok := existingStats[statKey]; ok && len(existing.Paths) > 0 { files = append(files, existing.Paths...) @@ -435,12 +413,12 @@ func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack) (map // Write batch stats blob via filesystem FFI. id, err := bw.allocator.AllocOne() if err != nil { - return nil, err + return nil, nil, err } relPath := fmt.Sprintf("_stats/bloom_filter.%d/%d", pkFieldID, id) fullPath := path.Join(basePath, relPath) if err := packed.WriteFile(bw.storageConfig, fullPath, batchStatsBlob.Value); err != nil { - return nil, err + return nil, nil, err } bw.sizeWritten += int64(len(batchStatsBlob.Value)) memorySize += int64(len(batchStatsBlob.Value)) @@ -452,54 +430,43 @@ func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack) (map // since the corresponding RollStats has been deferred above. mergedStatsBlob, err := serializer.serializeMergedPkStatsWith(pack, singlePKStats) if err != nil { - return nil, err + return nil, nil, err } mergedRelPath := fmt.Sprintf("_stats/bloom_filter.%d/%d", pkFieldID, int64(storage.CompoundStatsType)) mergedFullPath := path.Join(basePath, mergedRelPath) if err := packed.WriteFile(bw.storageConfig, mergedFullPath, mergedStatsBlob.Value); err != nil { - return nil, err + return nil, nil, err } bw.sizeWritten += int64(len(mergedStatsBlob.Value)) memorySize += int64(len(mergedStatsBlob.Value)) files = append(files, mergedFullPath) } - // Register stats in manifest - newManifest, err := packed.AddStatsToManifest(bw.manifestPath, bw.storageConfig, []packed.StatEntry{ - { - Key: statKey, - Files: files, - Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", memorySize)}, - }, - }) - if err != nil { - return nil, fmt.Errorf("failed to add stats to manifest: %w", err) - } - bw.manifestPath = newManifest + entries := []packed.StatEntry{{ + Key: statKey, + Files: files, + Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", memorySize)}, + }} - // Return empty map - stats are in manifest - return make(map[int64]*datapb.FieldBinlog), nil + // Return empty map - stats are embedded in the manifest, not separate binlogs + return make(map[int64]*datapb.FieldBinlog), entries, nil } -// writeBM25Stasts overrides the base class to write BM25 stats into the -// manifest instead of separate binlog files. -func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack) (map[int64]*datapb.FieldBinlog, error) { +// writeBM25Stasts writes BM25 stat blobs under basePath/_stats and returns +// the resulting StatEntry list. The caller folds the entries into a +// ManifestUpdates that commits atomically with inserts / delta / stats. +func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack, basePath string) (map[int64]*datapb.FieldBinlog, []packed.StatEntry, error) { if len(pack.bm25Stats) == 0 { - return make(map[int64]*datapb.FieldBinlog), nil + return make(map[int64]*datapb.FieldBinlog), nil, nil } serializer, err := NewStorageSerializer(bw.metaCache, bw.schema) if err != nil { - return nil, err + return nil, nil, err } bm25Blobs, err := serializer.serializeBM25Stats(pack) if err != nil { - return nil, err - } - - basePath, _, err := packed.UnmarshalManifestPath(bw.manifestPath) - if err != nil { - return nil, err + return nil, nil, err } // Track per-field files and memory sizes, then build stat entries at the end @@ -510,7 +477,7 @@ func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack) fieldMap := make(map[int64]*fieldStats) // Preserve existing BM25 stat files from previous batches. - existingStats, err := packed.GetManifestStats(bw.manifestPath, bw.storageConfig) + existingStats, err := packed.GetManifestStats(bw.initialManifestPath, bw.storageConfig) if err == nil { for key, existing := range existingStats { prefix, fieldID, ok := packed.ParseStatKey(key) @@ -528,13 +495,13 @@ func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack) for fieldID, blob := range bm25Blobs { id, err := bw.allocator.AllocOne() if err != nil { - return nil, err + return nil, nil, err } relPath := fmt.Sprintf("_stats/bm25.%d/%d", fieldID, id) fullPath := path.Join(basePath, relPath) if err := packed.WriteFile(bw.storageConfig, fullPath, blob.Value); err != nil { - return nil, err + return nil, nil, err } bw.sizeWritten += int64(len(blob.Value)) @@ -560,13 +527,13 @@ func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack) // since the corresponding MergeBm25Stats has been deferred above. mergedBM25Blob, err := serializer.serializeMergedBM25StatsWith(pack, pack.bm25Stats) if err != nil { - return nil, err + return nil, nil, err } for fieldID, blob := range mergedBM25Blob { mergedRelPath := fmt.Sprintf("_stats/bm25.%d/%d", fieldID, int64(storage.CompoundStatsType)) mergedFullPath := path.Join(basePath, mergedRelPath) if err := packed.WriteFile(bw.storageConfig, mergedFullPath, blob.Value); err != nil { - return nil, err + return nil, nil, err } bw.sizeWritten += int64(len(blob.Value)) @@ -580,26 +547,17 @@ func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack) } } - // Build stat entries with memory_size metadata - var statEntries []packed.StatEntry + var entries []packed.StatEntry for fieldID, fs := range fieldMap { - statEntries = append(statEntries, packed.StatEntry{ + entries = append(entries, packed.StatEntry{ Key: fmt.Sprintf("bm25.%d", fieldID), Files: fs.files, Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", fs.memorySize)}, }) } - if len(statEntries) > 0 { - newManifest, err := packed.AddStatsToManifest(bw.manifestPath, bw.storageConfig, statEntries) - if err != nil { - return nil, fmt.Errorf("failed to add BM25 stats to manifest: %w", err) - } - bw.manifestPath = newManifest - } - - // Return empty map - stats are in manifest - return make(map[int64]*datapb.FieldBinlog), nil + // Return empty map - stats are embedded in the manifest, not separate binlogs + return make(map[int64]*datapb.FieldBinlog), entries, nil } // buildTextColumnConfigs builds TextColumnConfig for all TEXT fields in the schema. diff --git a/internal/flushcommon/syncmgr/pack_writer_v3_helpers_test.go b/internal/flushcommon/syncmgr/pack_writer_v3_helpers_test.go new file mode 100644 index 0000000000..840737e71d --- /dev/null +++ b/internal/flushcommon/syncmgr/pack_writer_v3_helpers_test.go @@ -0,0 +1,87 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); 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 + +package syncmgr + +import ( + "fmt" + "testing" + + "github.com/cockroachdb/errors" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/milvus-io/milvus-proto/go-api/v3/schemapb" + "github.com/milvus-io/milvus/internal/storagev2/packed" + "github.com/milvus-io/milvus/pkg/v3/util/paramtable" + "github.com/milvus-io/milvus/pkg/v3/util/retry" +) + +func TestClassifyLoonErr(t *testing.T) { + t.Run("nil", func(t *testing.T) { + assert.NoError(t, classifyLoonErr(nil)) + }) + t.Run("transient stays retryable", func(t *testing.T) { + wrapped := fmt.Errorf("loon failed: %w", packed.ErrLoonTransient) + out := classifyLoonErr(wrapped) + require.Error(t, out) + assert.True(t, errors.Is(out, packed.ErrLoonTransient), + "transient error must remain retryable") + assert.True(t, retry.IsRecoverable(out), + "transient error must remain recoverable so retry.Do retries") + }) + t.Run("non-transient becomes unrecoverable", func(t *testing.T) { + out := classifyLoonErr(errors.New("permanent disk fault")) + require.Error(t, out) + assert.False(t, retry.IsRecoverable(out), + "non-transient error must be wrapped Unrecoverable so retry stops") + }) +} + +func TestBuildTextColumnConfigs(t *testing.T) { + // Required for paramtable.Get().DataNodeCfg.* to return defaults. + paramtable.Init() + + t.Run("no text fields returns nil", func(t *testing.T) { + schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{ + {FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true}, + {FieldID: 101, Name: "varchar", DataType: schemapb.DataType_VarChar}, + {FieldID: 102, Name: "vector", DataType: schemapb.DataType_FloatVector}, + }} + got := buildTextColumnConfigs(schema, "/seg/base") + assert.Empty(t, got, "schema without TEXT fields should produce no configs") + }) + + t.Run("one text field produces one config with derived lob path", func(t *testing.T) { + schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{ + {FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true}, + {FieldID: 101, Name: "doc", DataType: schemapb.DataType_Text}, + }} + got := buildTextColumnConfigs(schema, "files/insert_log/1/2") + require.Len(t, got, 1) + assert.Equal(t, int64(101), got[0].FieldID) + assert.Equal(t, "files/insert_log/1/2/lobs/101", got[0].LobBasePath) + }) + + t.Run("multiple text fields produce one config each", func(t *testing.T) { + schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{ + {FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true}, + {FieldID: 101, Name: "title", DataType: schemapb.DataType_Text}, + {FieldID: 102, Name: "body", DataType: schemapb.DataType_Text}, + }} + got := buildTextColumnConfigs(schema, "p") + require.Len(t, got, 2) + ids := []int64{got[0].FieldID, got[1].FieldID} + assert.ElementsMatch(t, []int64{101, 102}, ids) + for _, cfg := range got { + assert.Contains(t, cfg.LobBasePath, "lobs/") + } + }) +} diff --git a/internal/flushcommon/syncmgr/pack_writer_v3_test.go b/internal/flushcommon/syncmgr/pack_writer_v3_test.go index c254672c7c..5ce9a9efbf 100644 --- a/internal/flushcommon/syncmgr/pack_writer_v3_test.go +++ b/internal/flushcommon/syncmgr/pack_writer_v3_test.go @@ -24,6 +24,7 @@ import ( "sync/atomic" "testing" + "github.com/bytedance/mockey" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/suite" @@ -552,3 +553,112 @@ func (s *PackWriterV3Suite) TestMultiBatchBM25StatsAccumulation() { s.Greater(count2, count1, "batch 2 should accumulate more bm25 files than batch 1") s.NotEmpty(stats2[bm25Key].Metadata["memory_size"], "memory_size metadata should be set") } + +// TestWrite_SingleVersionBumpAcrossSections verifies that one Write call +// bumps the manifest version exactly once even when inserts, stats, delta, +// and bm25 are all present — the atomicity guarantee this refactor adds. +func (s *PackWriterV3Suite) TestWrite_SingleVersionBumpAcrossSections() { + collectionID := int64(123) + partitionID := int64(456) + segmentID := int64(789) + channelName := fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", collectionID) + rows := 10 + + bfs := pkoracle.NewBloomFilterSet() + k := metautil.JoinIDPath(collectionID, partitionID, segmentID) + basePath := path.Join(common.SegmentInsertLogPath, k) + manifestPath := packed.MarshalManifestPath(basePath, packed.ManifestEarliest) + + seg := metacache.NewSegmentInfo(&datapb.SegmentInfo{ManifestPath: manifestPath}, bfs, nil) + metacache.UpdateNumOfRows(1000)(seg) + mc := metacache.NewMockMetaCache(s.T()) + mc.EXPECT().Collection().Return(collectionID).Maybe() + mc.EXPECT().GetSchema(mock.Anything).Return(s.schema).Maybe() + mc.EXPECT().GetSegmentByID(segmentID).Return(seg, true).Maybe() + mc.EXPECT().GetSegmentsBy(mock.Anything, mock.Anything).Return([]*metacache.SegmentInfo{seg}).Maybe() + mc.EXPECT().UpdateSegments(mock.Anything, mock.Anything). + Run(func(action metacache.SegmentAction, filters ...metacache.SegmentFilter) { action(seg) }). + Return().Maybe() + + deletes := &storage.DeleteData{} + for i := 0; i < rows; i++ { + deletes.Append(storage.NewInt64PrimaryKey(int64(i+1)), uint64(100+i)) + } + + pack := new(SyncPack). + WithCollectionID(collectionID). + WithPartitionID(partitionID). + WithSegmentID(segmentID). + WithChannelName(channelName). + WithInsertData(genInsertData(rows, s.schema)). + WithDeleteData(deletes). + WithFlush() + + _, baseVer, err := packed.UnmarshalManifestPath(manifestPath) + s.Require().NoError(err) + + bw := NewBulkPackWriterV3(mc, s.schema, s.cm, s.logIDAlloc, + packed.DefaultWriteBufferSize, 0, s.storageConfig, s.currentSplit, manifestPath) + + _, _, _, _, writtenManifestPath, _, err := bw.Write(context.Background(), pack) + s.Require().NoError(err) + _, newVer, err := packed.UnmarshalManifestPath(writtenManifestPath) + s.Require().NoError(err) + s.Equalf(baseVer+1, newVer, + "Write must produce exactly one version bump; baseVer=%d newVer=%d", baseVer, newVer) +} + +// TestWrite_RetryDoesNotLeakVersionBumps verifies that when a transient +// commit failure forces the retry loop to re-run, the eventual successful +// commit produces only one version bump, not one per attempt. +func (s *PackWriterV3Suite) TestWrite_RetryDoesNotLeakVersionBumps() { + collectionID := int64(123) + partitionID := int64(456) + segmentID := int64(789) + channelName := fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", collectionID) + rows := 4 + + bfs := pkoracle.NewBloomFilterSet() + k := metautil.JoinIDPath(collectionID, partitionID, segmentID) + basePath := path.Join(common.SegmentInsertLogPath, k) + manifestPath := packed.MarshalManifestPath(basePath, packed.ManifestEarliest) + seg := metacache.NewSegmentInfo(&datapb.SegmentInfo{ManifestPath: manifestPath}, bfs, nil) + metacache.UpdateNumOfRows(1000)(seg) + mc := metacache.NewMockMetaCache(s.T()) + mc.EXPECT().Collection().Return(collectionID).Maybe() + mc.EXPECT().GetSchema(mock.Anything).Return(s.schema).Maybe() + mc.EXPECT().GetSegmentByID(segmentID).Return(seg, true).Maybe() + mc.EXPECT().GetSegmentsBy(mock.Anything, mock.Anything).Return([]*metacache.SegmentInfo{seg}).Maybe() + mc.EXPECT().UpdateSegments(mock.Anything, mock.Anything). + Run(func(action metacache.SegmentAction, filters ...metacache.SegmentFilter) { action(seg) }). + Return().Maybe() + + pack := new(SyncPack). + WithCollectionID(collectionID).WithPartitionID(partitionID).WithSegmentID(segmentID). + WithChannelName(channelName).WithInsertData(genInsertData(rows, s.schema)) + + _, baseVer, err := packed.UnmarshalManifestPath(manifestPath) + s.Require().NoError(err) + + var calls int32 + var origin func(string, int64, *indexpb.StorageConfig, *packed.ManifestUpdates) (string, error) + patched := mockey.Mock(packed.CommitManifestUpdates). + To(func(basePath string, baseVersion int64, cfg *indexpb.StorageConfig, updates *packed.ManifestUpdates) (string, error) { + if atomic.AddInt32(&calls, 1) == 1 { + return "", packed.ErrLoonTransient + } + return origin(basePath, baseVersion, cfg, updates) + }). + Origin(&origin). + Build() + defer patched.UnPatch() + + bw := NewBulkPackWriterV3(mc, s.schema, s.cm, s.logIDAlloc, + packed.DefaultWriteBufferSize, 0, s.storageConfig, s.currentSplit, manifestPath) + _, _, _, _, newManifestPath, _, err := bw.Write(context.Background(), pack) + s.Require().NoError(err) + _, newVer, err := packed.UnmarshalManifestPath(newManifestPath) + s.Require().NoError(err) + s.Equal(baseVer+1, newVer, "retry must not produce extra version bumps") + s.Equal(int32(2), atomic.LoadInt32(&calls), "exactly one retry expected") +} diff --git a/internal/storage/binlog_record_writer.go b/internal/storage/binlog_record_writer.go index 8c4aa6c4fc..cceb551337 100644 --- a/internal/storage/binlog_record_writer.go +++ b/internal/storage/binlog_record_writer.go @@ -66,6 +66,10 @@ type packedBinlogRecordWriterBase struct { columnGroups []storagecommon.ColumnGroup storageConfig *indexpb.StorageConfig storagePluginContext *indexcgopb.StoragePluginContext + // basePath is the segment data root, populated by initWriters. The + // underlying packed batch writers do not return a manifest path; the + // caller builds the manifest update against this base path. + basePath string pkCollector *PkStatsCollector bm25Collector *Bm25StatsCollector @@ -376,7 +380,7 @@ var _ BinlogRecordWriter = (*PackedManifestRecordWriter)(nil) type PackedManifestRecordWriter struct { packedBinlogRecordWriterBase // writer and stats generated at runtime - writer *packedRecordManifestWriter + writer *packedRecordBatchWriter } func (pw *PackedManifestRecordWriter) Write(r Record) error { @@ -424,8 +428,8 @@ func (pw *PackedManifestRecordWriter) initWriters(r Record) error { var err error k := metautil.JoinIDPath(pw.collectionID, pw.partitionID, pw.segmentID) - basePath := path.Join(pw.storageConfig.GetRootPath(), common.SegmentInsertLogPath, k) - pw.writer, err = NewPackedRecordManifestWriter(basePath, packed.ManifestEarliest, pw.schema, pw.bufferSize, pw.multiPartUploadSize, pw.columnGroups, pw.storageConfig, pw.storagePluginContext) + pw.basePath = path.Join(pw.storageConfig.GetRootPath(), common.SegmentInsertLogPath, k) + pw.writer, err = NewPackedRecordBatchWriter(pw.basePath, pw.schema, pw.bufferSize, pw.multiPartUploadSize, pw.columnGroups, pw.storageConfig, pw.storagePluginContext) if err != nil { return merr.WrapErrServiceInternal(fmt.Sprintf("can not new packed record writer %s", err.Error())) } @@ -459,38 +463,45 @@ func (pw *PackedManifestRecordWriter) finalizeBinlogs() { FieldNullCounts: pw.getFieldNullCountsForColumnGroup(columnGroup), }) } - pw.manifest = pw.writer.GetWrittenManifest() } +// Close finalizes the V3 segment. It closes the underlying packed writer to +// get column groups, serializes bloom filter / BM25 stat blobs, writes +// those blobs to storage, then performs a single packed.CommitManifestUpdates +// that registers inserts + all stats atomically. func (pw *PackedManifestRecordWriter) Close() error { - if pw.writer != nil { - if err := pw.writer.Close(); err != nil { - return err - } + if pw.writer == nil { + return nil + } + out, err := pw.writer.Close() + if err != nil { + return err + } + if out != nil { + defer out.Destroy() } pw.finalizeBinlogs() - // For V3 (manifest-based) segments, stats are stored under basePath/_stats/ - // and registered in the manifest. If pw.manifest is empty, no data was - // written, so there are no stats to write. - if pw.manifest != "" { - if err := pw.writeStatsV3(); err != nil { - return err - } + if out == nil { + return nil } + + updates := &packed.ManifestUpdates{NewFiles: out} + if err := pw.appendV3Stats(updates); err != nil { + return err + } + newManifest, err := packed.CommitManifestUpdates(pw.basePath, packed.ManifestEarliest, pw.storageConfig, updates) + if err != nil { + return fmt.Errorf("PackedManifestRecordWriter.Close commit: %w", err) + } + pw.manifest = newManifest return nil } -// writeStatsV3 writes bloom filter and BM25 stats for V3 (manifest-based) -// segments. Stats blobs are written to basePath/_stats/ and registered in the -// manifest via a transaction, leaving pw.statsLog and pw.bm25StatsLog nil so -// callers know stats are embedded in the manifest. -func (pw *PackedManifestRecordWriter) writeStatsV3() error { - basePath, _, err := packed.UnmarshalManifestPath(pw.manifest) - if err != nil { - return fmt.Errorf("writeStatsV3: failed to parse manifest path: %w", err) - } - - // --- Bloom filter (PK) stats --- +// appendV3Stats serializes bloom filter / BM25 stat blobs, writes them to +// storage, and appends StatEntry records onto updates so the surrounding +// commit registers inserts + stats atomically. Leaves pw.statsLog and +// pw.bm25StatsLog nil so callers know stats are embedded in the manifest. +func (pw *PackedManifestRecordWriter) appendV3Stats(updates *packed.ManifestUpdates) error { statsBlob, pkFieldID, err := pw.pkCollector.SerializeBlob(pw.rowNum) if err != nil { return err @@ -500,55 +511,36 @@ func (pw *PackedManifestRecordWriter) writeStatsV3() error { if err != nil { return err } - fullPath := path.Join(basePath, fmt.Sprintf("_stats/bloom_filter.%d/%d", pkFieldID, id)) + fullPath := path.Join(pw.basePath, fmt.Sprintf("_stats/bloom_filter.%d/%d", pkFieldID, id)) if err := packed.WriteFile(pw.storageConfig, fullPath, statsBlob.Value); err != nil { - return fmt.Errorf("writeStatsV3: failed to write bloom filter stats: %w", err) + return fmt.Errorf("appendV3Stats: failed to write bloom filter stats: %w", err) } - - newManifest, err := packed.AddStatsToManifest(pw.manifest, pw.storageConfig, []packed.StatEntry{ - { - Key: fmt.Sprintf("bloom_filter.%d", pkFieldID), - Files: []string{fullPath}, - Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", int64(len(statsBlob.Value)))}, - }, + updates.Stats = append(updates.Stats, packed.StatEntry{ + Key: fmt.Sprintf("bloom_filter.%d", pkFieldID), + Files: []string{fullPath}, + Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", int64(len(statsBlob.Value)))}, }) - if err != nil { - return fmt.Errorf("writeStatsV3: failed to add bloom filter stats to manifest: %w", err) - } - pw.manifest = newManifest - // Stats are stored in the manifest; leave pw.statsLog nil. } - // --- BM25 stats --- bm25Blobs, err := pw.bm25Collector.SerializeBlobs() if err != nil { return err } - if len(bm25Blobs) > 0 { - var statEntries []packed.StatEntry - for fieldID, blob := range bm25Blobs { - id, err := pw.allocator.AllocOne() - if err != nil { - return err - } - fullPath := path.Join(basePath, fmt.Sprintf("_stats/bm25.%d/%d", fieldID, id)) - if err := packed.WriteFile(pw.storageConfig, fullPath, blob.Value); err != nil { - return fmt.Errorf("writeStatsV3: failed to write bm25 stats: %w", err) - } - statEntries = append(statEntries, packed.StatEntry{ - Key: fmt.Sprintf("bm25.%d", fieldID), - Files: []string{fullPath}, - Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", blob.MemorySize)}, - }) - } - newManifest, err := packed.AddStatsToManifest(pw.manifest, pw.storageConfig, statEntries) + for fieldID, blob := range bm25Blobs { + id, err := pw.allocator.AllocOne() if err != nil { - return fmt.Errorf("writeStatsV3: failed to add bm25 stats to manifest: %w", err) + return err } - pw.manifest = newManifest - // Stats are stored in the manifest; leave pw.bm25StatsLog nil. + fullPath := path.Join(pw.basePath, fmt.Sprintf("_stats/bm25.%d/%d", fieldID, id)) + if err := packed.WriteFile(pw.storageConfig, fullPath, blob.Value); err != nil { + return fmt.Errorf("appendV3Stats: failed to write bm25 stats: %w", err) + } + updates.Stats = append(updates.Stats, packed.StatEntry{ + Key: fmt.Sprintf("bm25.%d", fieldID), + Files: []string{fullPath}, + Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", blob.MemorySize)}, + }) } - return nil } @@ -601,11 +593,11 @@ func newPackedManifestRecordWriter(collectionID, partitionID, segmentID UniqueID var _ BinlogRecordWriter = (*PackedTextManifestRecordWriter)(nil) -// PackedTextManifestRecordWriter wraps packedTextManifestWriter for TEXT column compaction. +// PackedTextManifestRecordWriter wraps packedTextBatchWriter for TEXT column compaction. // this writer is used during compaction when TEXT columns need REWRITE_ALL strategy. type PackedTextManifestRecordWriter struct { packedBinlogRecordWriterBase - writer *packedTextManifestWriter + writer *packedTextBatchWriter textColumnConfigs []packed.TextColumnConfig } @@ -654,8 +646,8 @@ func (pw *PackedTextManifestRecordWriter) initWriters(r Record) error { var err error k := metautil.JoinIDPath(pw.collectionID, pw.partitionID, pw.segmentID) - basePath := path.Join(pw.storageConfig.GetRootPath(), common.SegmentInsertLogPath, k) - pw.writer, err = NewPackedTextManifestWriter(pw.storageConfig.GetBucketName(), basePath, -1, pw.schema, pw.bufferSize, pw.multiPartUploadSize, pw.columnGroups, pw.storageConfig, pw.textColumnConfigs) + pw.basePath = path.Join(pw.storageConfig.GetRootPath(), common.SegmentInsertLogPath, k) + pw.writer, err = NewPackedTextBatchWriter(pw.storageConfig.GetBucketName(), pw.basePath, pw.schema, pw.bufferSize, pw.multiPartUploadSize, pw.columnGroups, pw.storageConfig, pw.textColumnConfigs) if err != nil { return merr.WrapErrServiceInternal(fmt.Sprintf("can not new packed text writer %s", err.Error())) } @@ -689,20 +681,33 @@ func (pw *PackedTextManifestRecordWriter) finalizeBinlogs() { FieldNullCounts: pw.getFieldNullCountsForColumnGroup(columnGroup), }) } - pw.manifest = pw.writer.GetWrittenManifest() } +// Close finalizes the text-column segment using the same do-then-commit +// pattern as PackedManifestRecordWriter.Close. func (pw *PackedTextManifestRecordWriter) Close() error { - if pw.writer != nil { - if err := pw.writer.Close(); err != nil { - return err - } + if pw.writer == nil { + return pw.writeStats() } - pw.finalizeBinlogs() - if err := pw.writeStats(); err != nil { + out, err := pw.writer.Close() + if err != nil { return err } - return nil + if out != nil { + defer out.Destroy() + } + pw.finalizeBinlogs() + if out == nil { + return pw.writeStats() + } + + newManifest, err := packed.CommitManifestUpdates(pw.basePath, packed.ManifestEarliest, pw.storageConfig, + &packed.ManifestUpdates{NewFiles: out}) + if err != nil { + return fmt.Errorf("PackedTextManifestRecordWriter.Close commit: %w", err) + } + pw.manifest = newManifest + return pw.writeStats() } // NewPackedTextManifestRecordWriter creates a new BinlogRecordWriter for TEXT column compaction. diff --git a/internal/storage/binlog_record_writer_v3_test.go b/internal/storage/binlog_record_writer_v3_test.go new file mode 100644 index 0000000000..67058f6c2a --- /dev/null +++ b/internal/storage/binlog_record_writer_v3_test.go @@ -0,0 +1,77 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); 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 + +package storage + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/milvus-io/milvus-proto/go-api/v3/schemapb" + "github.com/milvus-io/milvus/internal/allocator" + "github.com/milvus-io/milvus/pkg/v3/common" + "github.com/milvus-io/milvus/pkg/v3/proto/indexpb" +) + +// TestPackedManifestRecordWriter_CloseWithoutWrite verifies the no-data +// short-circuit: Close on a freshly-constructed PackedManifestRecordWriter +// must not call into the FFI commit path and must not produce a manifest. +func TestPackedManifestRecordWriter_CloseWithoutWrite(t *testing.T) { + schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{ + {FieldID: common.TimeStampField, DataType: schemapb.DataType_Int64}, + {FieldID: common.RowIDField, DataType: schemapb.DataType_Int64, IsPrimaryKey: true}, + }} + + dir := t.TempDir() + cfg := &indexpb.StorageConfig{StorageType: "local", RootPath: dir} + + w, err := newPackedManifestRecordWriter(1, 2, 3, schema, + ChunkedBlobsWriter(func(_ []*Blob) error { return nil }), + allocator.NewLocalAllocator(1, 1<<20), + 1024, 0, 0, nil, cfg, nil) + require.NoError(t, err) + + // No Write before Close. The internal `writer` field stays nil so + // Close must return immediately without invoking the FFI commit. + require.NoError(t, w.Close()) + + _, statsLog, _, manifestPath, _ := w.GetLogs() + assert.Empty(t, manifestPath, "no Write means no manifest should be produced") + assert.Nil(t, statsLog, "no Write means no statsLog should be produced") +} + +// TestPackedTextManifestRecordWriter_CloseWithoutWrite is the parallel +// short-circuit test for the text writer: Close with no Writes must +// still succeed without touching the FFI segment writer. +func TestPackedTextManifestRecordWriter_CloseWithoutWrite(t *testing.T) { + schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{ + {FieldID: common.TimeStampField, DataType: schemapb.DataType_Int64}, + {FieldID: common.RowIDField, DataType: schemapb.DataType_Int64, IsPrimaryKey: true}, + {FieldID: 101, DataType: schemapb.DataType_Text, Name: "doc"}, + }} + + dir := t.TempDir() + cfg := &indexpb.StorageConfig{StorageType: "local", RootPath: dir} + + w, err := NewPackedTextManifestRecordWriter(1, 2, 3, schema, + ChunkedBlobsWriter(func(_ []*Blob) error { return nil }), + allocator.NewLocalAllocator(1, 1<<20), + 1024, 0, 0, nil, cfg, nil) + require.NoError(t, err) + + // No Write before Close. The text writer's nil-handling path must + // fall through to writeStats() and exit cleanly. + require.NoError(t, w.Close()) + + _, _, _, manifestPath, _ := w.GetLogs() + assert.Empty(t, manifestPath, "no Write means no manifest should be produced") +} diff --git a/internal/storage/record_writer.go b/internal/storage/record_writer.go index 578b5ecfb8..764531582c 100644 --- a/internal/storage/record_writer.go +++ b/internal/storage/record_writer.go @@ -180,7 +180,7 @@ func NewPackedRecordWriter( }, nil } -type packedRecordManifestWriter struct { +type packedRecordBatchWriter struct { writer *packed.FFIPackedWriter bufferSize int64 columnGroups []storagecommon.ColumnGroup @@ -192,11 +192,10 @@ type packedRecordManifestWriter struct { writtenUncompressed uint64 columnGroupUncompressed map[typeutil.UniqueID]uint64 columnGroupCompressed map[typeutil.UniqueID]uint64 - outputManifest string storageConfig *indexpb.StorageConfig } -func (pw *packedRecordManifestWriter) Write(r Record) error { +func (pw *packedRecordBatchWriter) Write(r Record) error { var rec arrow.Record sar, ok := r.(*simpleArrowRecord) if !ok { @@ -225,63 +224,63 @@ func (pw *packedRecordManifestWriter) Write(r Record) error { return pw.writer.WriteRecordBatch(rec) } -func (pw *packedRecordManifestWriter) GetWrittenUncompressed() uint64 { +func (pw *packedRecordBatchWriter) GetWrittenUncompressed() uint64 { return pw.writtenUncompressed } -func (pw *packedRecordManifestWriter) GetColumnGroupWrittenUncompressed(columnGroup typeutil.UniqueID) uint64 { +func (pw *packedRecordBatchWriter) GetColumnGroupWrittenUncompressed(columnGroup typeutil.UniqueID) uint64 { if size, ok := pw.columnGroupUncompressed[columnGroup]; ok { return size } return 0 } -func (pw *packedRecordManifestWriter) GetColumnGroupWrittenCompressed(columnGroup typeutil.UniqueID) uint64 { +func (pw *packedRecordBatchWriter) GetColumnGroupWrittenCompressed(columnGroup typeutil.UniqueID) uint64 { if size, ok := pw.columnGroupCompressed[columnGroup]; ok { return size } return 0 } -func (pw *packedRecordManifestWriter) GetWrittenPaths(columnGroup typeutil.UniqueID) string { +func (pw *packedRecordBatchWriter) GetWrittenPaths(columnGroup typeutil.UniqueID) string { if path, ok := pw.pathsMap[columnGroup]; ok { return path } return "" } -func (pw *packedRecordManifestWriter) GetWrittenManifest() string { - return pw.outputManifest -} - -func (pw *packedRecordManifestWriter) GetWrittenRowNum() int64 { +func (pw *packedRecordBatchWriter) GetWrittenRowNum() int64 { return pw.rowNum } -func (pw *packedRecordManifestWriter) Close() error { - if pw.writer != nil { - manifest, err := pw.writer.Close() - if err != nil { - return err - } - pw.outputManifest = manifest - for id := range pw.pathsMap { - pw.columnGroupCompressed[id] = uint64(0) - } +// Close closes the underlying FFI writer and returns the resulting +// column-groups payload. The writer never touches the manifest — the +// caller passes the returned handle to packed.CommitManifestUpdates and +// calls Destroy on it when done. +func (pw *packedRecordBatchWriter) Close() (packed.WriterOutput, error) { + if pw.writer == nil { + return nil, nil } - return nil + out, err := pw.writer.Close() + if err != nil { + return nil, err + } + pw.writer = nil + for id := range pw.pathsMap { + pw.columnGroupCompressed[id] = uint64(0) + } + return out, nil } -func NewPackedRecordManifestWriter( +func NewPackedRecordBatchWriter( basePath string, - baseVersion int64, schema *schemapb.CollectionSchema, bufferSize int64, multiPartUploadSize int64, columnGroups []storagecommon.ColumnGroup, storageConfig *indexpb.StorageConfig, storagePluginContext *indexcgopb.StoragePluginContext, -) (*packedRecordManifestWriter, error) { +) (*packedRecordBatchWriter, error) { // Validate PK field exists before proceeding _, err := typeutil.GetPrimaryFieldSchema(schema) if err != nil { @@ -294,7 +293,7 @@ func NewPackedRecordManifestWriter( fmt.Sprintf("can not convert collection schema %s to arrow schema: %s", schema.Name, err.Error())) } - writer, err := packed.NewFFIPackedWriter(basePath, baseVersion, arrowSchema, columnGroups, storageConfig, storagePluginContext) + writer, err := packed.NewFFIPackedWriter(basePath, arrowSchema, columnGroups, storageConfig, storagePluginContext) if err != nil { return nil, merr.WrapErrServiceInternal( fmt.Sprintf("can not new packed record writer %s", err.Error())) @@ -312,7 +311,7 @@ func NewPackedRecordManifestWriter( pathsMap[columnGroup.GroupID] = path.Join(basePath, strconv.FormatInt(columnGroup.GroupID, 10), strconv.FormatInt(start, 10)) } - return &packedRecordManifestWriter{ + return &packedRecordBatchWriter{ writer: writer, schema: schema, arrowSchema: arrowSchema, @@ -339,9 +338,9 @@ func NewPackedSerializeWriter(bucketName string, paths []string, schema *schemap }, batchSize), nil } -// packedTextManifestWriter wraps FFISegmentWriter for TEXT column support during compaction. +// packedTextBatchWriter wraps FFISegmentWriter for TEXT column support during compaction. // it handles TEXT column rewriting with LOB file management in REWRITE_ALL mode. -type packedTextManifestWriter struct { +type packedTextBatchWriter struct { writer *packed.FFISegmentWriter bufferSize int64 columnGroups []storagecommon.ColumnGroup @@ -353,11 +352,10 @@ type packedTextManifestWriter struct { writtenUncompressed uint64 columnGroupUncompressed map[typeutil.UniqueID]uint64 columnGroupCompressed map[typeutil.UniqueID]uint64 - outputManifest string storageConfig *indexpb.StorageConfig } -func (pw *packedTextManifestWriter) Write(r Record) error { +func (pw *packedTextBatchWriter) Write(r Record) error { var rec arrow.Record sar, ok := r.(*simpleArrowRecord) if !ok { @@ -386,69 +384,71 @@ func (pw *packedTextManifestWriter) Write(r Record) error { return pw.writer.Write(rec) } -func (pw *packedTextManifestWriter) GetWrittenUncompressed() uint64 { +func (pw *packedTextBatchWriter) GetWrittenUncompressed() uint64 { return pw.writtenUncompressed } -func (pw *packedTextManifestWriter) GetColumnGroupWrittenUncompressed(columnGroup typeutil.UniqueID) uint64 { +func (pw *packedTextBatchWriter) GetColumnGroupWrittenUncompressed(columnGroup typeutil.UniqueID) uint64 { if size, ok := pw.columnGroupUncompressed[columnGroup]; ok { return size } return 0 } -func (pw *packedTextManifestWriter) GetColumnGroupWrittenCompressed(columnGroup typeutil.UniqueID) uint64 { +func (pw *packedTextBatchWriter) GetColumnGroupWrittenCompressed(columnGroup typeutil.UniqueID) uint64 { if size, ok := pw.columnGroupCompressed[columnGroup]; ok { return size } return 0 } -func (pw *packedTextManifestWriter) GetWrittenPaths(columnGroup typeutil.UniqueID) string { +func (pw *packedTextBatchWriter) GetWrittenPaths(columnGroup typeutil.UniqueID) string { if path, ok := pw.pathsMap[columnGroup]; ok { return path } return "" } -func (pw *packedTextManifestWriter) GetWrittenManifest() string { - return pw.outputManifest -} - -func (pw *packedTextManifestWriter) GetWrittenRowNum() int64 { +func (pw *packedTextBatchWriter) GetWrittenRowNum() int64 { return pw.rowNum } -func (pw *packedTextManifestWriter) Close() error { - if pw.writer != nil { - defer pw.writer.Destroy() - result, err := pw.writer.Close() - if err != nil { - return err - } - pw.outputManifest = result.ManifestPath - pw.rowNum = result.RowsWritten - for id := range pw.pathsMap { - pw.columnGroupCompressed[id] = uint64(0) - } +// Close closes the underlying segment writer and returns the resulting +// column-groups + LOB payload. The writer never touches the manifest — +// the caller passes the returned handle to packed.CommitManifestUpdates +// and calls Destroy on it when done. FFISegmentWriter.Close releases its +// own handle and properties, so no extra cleanup is needed here. +func (pw *packedTextBatchWriter) Close() (packed.WriterOutput, error) { + if pw.writer == nil { + return nil, nil } - return nil + out, err := pw.writer.Close() + pw.writer = nil + if err != nil { + return nil, err + } + if so, ok := out.(*packed.SegmentOutput); ok { + pw.rowNum = so.RowsWritten() + } + for id := range pw.pathsMap { + pw.columnGroupCompressed[id] = uint64(0) + } + return out, nil } -// NewPackedTextManifestWriter creates a new writer that uses FFISegmentWriter with TEXT column support. +// NewPackedTextBatchWriter creates a new writer that uses FFISegmentWriter with TEXT column support. // this writer is used during compaction when TEXT columns need REWRITE_ALL strategy. // textColumnConfigs: TEXT column configurations for REWRITE_ALL fields (nil if no TEXT fields need rewriting) -func NewPackedTextManifestWriter( +func NewPackedTextBatchWriter( bucketName string, basePath string, - baseVersion int64, schema *schemapb.CollectionSchema, bufferSize int64, multiPartUploadSize int64, columnGroups []storagecommon.ColumnGroup, storageConfig *indexpb.StorageConfig, textColumnConfigs []packed.TextColumnConfig, -) (*packedTextManifestWriter, error) { +) (*packedTextBatchWriter, error) { // validate PK field exists before proceeding _, err := typeutil.GetPrimaryFieldSchema(schema) if err != nil { @@ -477,8 +477,6 @@ func NewPackedTextManifestWriter( // build segment writer config config := &packed.SegmentWriterConfig{ SegmentPath: basePath, - ReadVersion: baseVersion, - RetryLimit: 3, TextColumns: textColumnConfigs, } @@ -501,7 +499,7 @@ func NewPackedTextManifestWriter( pathsMap[columnGroup.GroupID] = path.Join(basePath, strconv.FormatInt(columnGroup.GroupID, 10), strconv.FormatInt(start, 10)) } - return &packedTextManifestWriter{ + return &packedTextBatchWriter{ writer: writer, schema: schema, arrowSchema: arrowSchema, diff --git a/internal/storagev2/packed/manifest_commit.go b/internal/storagev2/packed/manifest_commit.go new file mode 100644 index 0000000000..f2b4dbb22d --- /dev/null +++ b/internal/storagev2/packed/manifest_commit.go @@ -0,0 +1,158 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); 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 + +package packed + +/* +#cgo pkg-config: milvus_core milvus-storage + +#include +#include "milvus-storage/ffi_c.h" +*/ +import "C" + +import ( + "fmt" + "unsafe" + + "github.com/milvus-io/milvus/pkg/v3/proto/indexpb" +) + +// WriterOutput is the data carrier returned by an FFI writer's Close. It +// owns C memory and must be released via Destroy after the surrounding +// CommitManifestUpdates call returns (success or failure). +// +// Concrete implementations live alongside the writers that produce them: +// - *ColumnGroups (FFIPackedWriter.Close) +// - *SegmentOutput (FFISegmentWriter.Close) +// +// Each implementation knows how to stage its payload onto a loon +// transaction handle via the package-internal applyTo method. +type WriterOutput interface { + // Destroy releases the underlying C resources. Idempotent. + Destroy() + // applyTo stages the output onto a loon transaction handle. Called by + // applyManifestUpdates as part of CommitManifestUpdates. + applyTo(handle C.LoonTransactionHandle) error +} + +// ManifestUpdates bundles every data-file-level change a single caller +// wants to apply atomically to a manifest. The slow file writes (parquet, +// deltalog, stat blobs) happen before this payload is assembled; +// CommitManifestUpdates then opens a transaction, applies everything in +// one shot, and commits. +// +// NewFiles holds C memory produced by an FFI writer and MUST be released +// by the caller via Destroy after CommitManifestUpdates returns (success +// or failure). +type ManifestUpdates struct { + // NewFiles is the column-groups / LOB payload returned by an FFI + // writer's Close. nil if no insert files were written. + NewFiles WriterOutput + // DeltaLogs is the list of delta-log entries to register. + DeltaLogs []DeltaLogEntry + // Stats is the list of stat entries (bloom filter, bm25, etc.) to + // register. Each entry's Files / Metadata replace any existing entry + // with the same Key (loon overwrite semantics). + Stats []StatEntry +} + +// isEmpty short-circuits CommitManifestUpdates when the caller assembled +// no work. A non-nil NewFiles is treated as real work; callers must not +// pass an already-destroyed payload here (the C transaction would then +// commit with zero staged ops and the loon side would return +// "Cannot commit: no updates recorded"). +func (u *ManifestUpdates) isEmpty() bool { + if u == nil { + return true + } + if u.NewFiles != nil { + return false + } + return len(u.DeltaLogs) == 0 && len(u.Stats) == 0 +} + +// CommitManifestUpdates opens a loon transaction at (basePath, baseVersion), +// applies every change in updates, commits, and returns the new manifest +// path. With no effective changes it returns the unchanged manifest path +// without opening a transaction. +// +// The caller retains ownership of any WriterOutput referenced by updates +// and is responsible for calling Destroy on it after CommitManifestUpdates +// returns. +func CommitManifestUpdates(basePath string, baseVersion int64, + storageConfig *indexpb.StorageConfig, updates *ManifestUpdates, +) (string, error) { + if updates.isEmpty() { + return MarshalManifestPath(basePath, baseVersion), nil + } + + cProperties, err := MakePropertiesFromStorageConfig(storageConfig, nil) + if err != nil { + return "", fmt.Errorf("commit manifest: %w", err) + } + defer C.loon_properties_free(cProperties) + + cBasePath := C.CString(basePath) + defer C.free(unsafe.Pointer(cBasePath)) + + // Use OVERWRITE for the whole bundle: stats updates rely on + // overwrite-on-key semantics (loon_transaction_update_stat replaces + // existing entries with the same key), while column-group appends and + // delta-log adds don't collide on a key in the first place. A single + // resolve mode keeps the API simple and matches the prior + // AddStatsToManifest behavior. + var handle C.LoonTransactionHandle + res := C.loon_transaction_begin(cBasePath, cProperties, + C.int64_t(baseVersion), + C.LOON_TRANSACTION_RESOLVE_OVERWRITE, + getRetryLimit(), &handle) + if err := HandleLoonFFIResult(res); err != nil { + return "", fmt.Errorf("commit manifest begin: %w", err) + } + defer C.loon_transaction_destroy(handle) + + if err := applyManifestUpdates(handle, updates); err != nil { + return "", err + } + + var commitVersion C.int64_t + res = C.loon_transaction_commit(handle, &commitVersion) + if err := HandleLoonFFIResult(res); err != nil { + return "", fmt.Errorf("commit manifest commit: %w", err) + } + return MarshalManifestPath(basePath, int64(commitVersion)), nil +} + +// applyManifestUpdates stages every operation in updates onto the loon +// transaction handle. +func applyManifestUpdates(handle C.LoonTransactionHandle, updates *ManifestUpdates) error { + if updates.NewFiles != nil { + if err := updates.NewFiles.applyTo(handle); err != nil { + return err + } + } + + for _, entry := range updates.DeltaLogs { + cPath := C.CString(entry.Path) + err := HandleLoonFFIResult(C.loon_transaction_add_delta_log(handle, cPath, C.int64_t(entry.NumEntries))) + C.free(unsafe.Pointer(cPath)) + if err != nil { + return fmt.Errorf("commit manifest add_delta_log: %w", err) + } + } + + for _, entry := range updates.Stats { + if err := UpdateTransactionStat(handle, entry.Key, entry.Files, entry.Metadata); err != nil { + return fmt.Errorf("commit manifest update_stat: %w", err) + } + } + return nil +} diff --git a/internal/storagev2/packed/manifest_commit_test.go b/internal/storagev2/packed/manifest_commit_test.go new file mode 100644 index 0000000000..468426c581 --- /dev/null +++ b/internal/storagev2/packed/manifest_commit_test.go @@ -0,0 +1,295 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); 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 + +package packed + +import ( + "path" + "testing" + + "github.com/apache/arrow/go/v17/arrow" + "github.com/apache/arrow/go/v17/arrow/array" + "github.com/apache/arrow/go/v17/arrow/memory" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/milvus-io/milvus/internal/storagecommon" + "github.com/milvus-io/milvus/pkg/v3/proto/indexpb" + "github.com/milvus-io/milvus/pkg/v3/util/paramtable" +) + +// manifestTestStorageConfig returns a storage config wired to a per-test +// temp dir for use with CommitManifestUpdates / FFIPackedWriter. +func manifestTestStorageConfig(t *testing.T) *indexpb.StorageConfig { + t.Helper() + paramtable.Init() + pt := paramtable.Get() + pt.Save(pt.CommonCfg.StorageType.Key, "local") + dir := t.TempDir() + pt.Save(pt.LocalStorageCfg.Path.Key, dir) + t.Cleanup(func() { + pt.Reset(pt.CommonCfg.StorageType.Key) + pt.Reset(pt.LocalStorageCfg.Path.Key) + }) + return &indexpb.StorageConfig{StorageType: "local", RootPath: dir} +} + +// TestCommitManifestUpdates_EmptyShortCircuit verifies that an empty +// ManifestUpdates returns the unchanged manifest path without opening a +// loon transaction (and therefore without bumping the version). +func TestCommitManifestUpdates_EmptyShortCircuit(t *testing.T) { + cfg := manifestTestStorageConfig(t) + basePath := "files/commit_empty_test/seg1" + + cases := []struct { + name string + u *ManifestUpdates + }{ + {"nil-updates", nil}, + {"zero-updates", &ManifestUpdates{}}, + {"empty-slices", &ManifestUpdates{DeltaLogs: []DeltaLogEntry{}, Stats: []StatEntry{}}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg, tc.u) + require.NoError(t, err) + require.Equal(t, MarshalManifestPath(basePath, ManifestEarliest), got, + "empty updates must return the unchanged manifest path") + }) + } +} + +// TestCommitManifestUpdates_StatsOnly exercises the stats-without-inserts +// path: a stat blob is written to storage and registered in the manifest +// via a single CommitManifestUpdates call. +func TestCommitManifestUpdates_StatsOnly(t *testing.T) { + cfg := manifestTestStorageConfig(t) + basePath := "files/commit_stats_only/seg1" + + statPath := path.Join(cfg.RootPath, basePath, "_stats/bloom_filter.100/1") + require.NoError(t, WriteFile(cfg, statPath, []byte("bloom-blob"))) + + got, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg, &ManifestUpdates{ + Stats: []StatEntry{{ + Key: "bloom_filter.100", + Files: []string{statPath}, + Metadata: map[string]string{"memory_size": "10"}, + }}, + }) + require.NoError(t, err) + _, version, err := UnmarshalManifestPath(got) + require.NoError(t, err) + require.Equal(t, int64(1), version, "stats-only commit must bump version exactly once") + + stats, err := GetManifestStats(got, cfg) + require.NoError(t, err) + require.Contains(t, stats, "bloom_filter.100") + require.Equal(t, "10", stats["bloom_filter.100"].Metadata["memory_size"]) +} + +// TestCommitManifestUpdates_DeltaOnly exercises the delta-only path. +func TestCommitManifestUpdates_DeltaOnly(t *testing.T) { + cfg := manifestTestStorageConfig(t) + basePath := "files/commit_delta_only/seg1" + + deltaPath := path.Join(cfg.RootPath, basePath, "_delta/delta-1") + require.NoError(t, WriteFile(cfg, deltaPath, []byte("delta-payload"))) + + got, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg, &ManifestUpdates{ + DeltaLogs: []DeltaLogEntry{{Path: deltaPath, NumEntries: 7}}, + }) + require.NoError(t, err) + _, version, err := UnmarshalManifestPath(got) + require.NoError(t, err) + require.Equal(t, int64(1), version) + + paths, err := GetDeltaLogPathsFromManifest(got, cfg) + require.NoError(t, err) + require.Len(t, paths, 1) +} + +// TestCommitManifestUpdates_AllSections is the integration variant: one +// commit covering inserts (column-groups from a real FFIPackedWriter), +// stats, and delta entries — asserts a single version bump for the whole +// bundle. +func TestCommitManifestUpdates_AllSections(t *testing.T) { + cfg := manifestTestStorageConfig(t) + basePath := "files/commit_all_sections/seg1" + + schema := arrow.NewSchema([]arrow.Field{ + { + Name: "pk", + Type: arrow.PrimitiveTypes.Int64, + Nullable: false, + Metadata: arrow.NewMetadata([]string{ArrowFieldIdMetadataKey}, []string{"100"}), + }, + }, nil) + b := array.NewRecordBuilder(memory.DefaultAllocator, schema) + defer b.Release() + pkb := b.Field(0).(*array.Int64Builder) + for i := 0; i < 4; i++ { + pkb.Append(int64(i)) + } + rec := b.NewRecord() + defer rec.Release() + + columnGroups := []storagecommon.ColumnGroup{ + {Columns: []int{0}, GroupID: storagecommon.DefaultShortColumnGroupID}, + } + + w, err := NewFFIPackedWriter(basePath, schema, columnGroups, cfg, nil) + require.NoError(t, err) + require.NoError(t, w.WriteRecordBatch(rec)) + out, err := w.Close() + require.NoError(t, err) + defer out.Destroy() + + statPath := path.Join(cfg.RootPath, basePath, "_stats/bloom_filter.100/1") + require.NoError(t, WriteFile(cfg, statPath, []byte("bloom-blob"))) + deltaPath := path.Join(cfg.RootPath, basePath, "_delta/delta-1") + require.NoError(t, WriteFile(cfg, deltaPath, []byte("delta-payload"))) + + got, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg, &ManifestUpdates{ + NewFiles: out, + DeltaLogs: []DeltaLogEntry{{Path: deltaPath, NumEntries: 4}}, + Stats: []StatEntry{{ + Key: "bloom_filter.100", + Files: []string{statPath}, + Metadata: map[string]string{"memory_size": "10"}, + }}, + }) + require.NoError(t, err) + + _, version, err := UnmarshalManifestPath(got) + require.NoError(t, err) + require.Equal(t, int64(1), version, + "inserts + delta + stats together must bump version exactly once") +} + +// TestCommitManifestUpdates_AddNewColumnGroups exercises the +// add_column_group (function-backfill) branch in applyManifestUpdates. +func TestCommitManifestUpdates_AddNewColumnGroups(t *testing.T) { + cfg := manifestTestStorageConfig(t) + basePath := "files/commit_new_cgs/seg1" + + schema := arrow.NewSchema([]arrow.Field{ + { + Name: "pk", + Type: arrow.PrimitiveTypes.Int64, + Nullable: false, + Metadata: arrow.NewMetadata([]string{ArrowFieldIdMetadataKey}, []string{"100"}), + }, + }, nil) + b := array.NewRecordBuilder(memory.DefaultAllocator, schema) + defer b.Release() + b.Field(0).(*array.Int64Builder).Append(int64(1)) + rec := b.NewRecord() + defer rec.Release() + + columnGroups := []storagecommon.ColumnGroup{ + {Columns: []int{0}, GroupID: storagecommon.DefaultShortColumnGroupID}, + } + + w, err := NewFFIPackedWriter(basePath, schema, columnGroups, cfg, nil) + require.NoError(t, err) + w.AsNewColumnGroups() + require.NoError(t, w.WriteRecordBatch(rec)) + out, err := w.Close() + require.NoError(t, err) + defer out.Destroy() + cgs, ok := out.(*ColumnGroups) + require.True(t, ok, "Close should return *ColumnGroups for FFIPackedWriter") + require.True(t, cgs.addNewColumnGroups, + "AsNewColumnGroups should propagate into ColumnGroups handle") + + got, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg, + &ManifestUpdates{NewFiles: out}) + require.NoError(t, err) + _, version, err := UnmarshalManifestPath(got) + require.NoError(t, err) + require.Equal(t, int64(1), version) +} + +// TestColumnGroups_DestroyIdempotent verifies that calling Destroy on a +// ColumnGroups handle multiple times (and on a nil receiver) is safe. +func TestColumnGroups_DestroyIdempotent(t *testing.T) { + cfg := manifestTestStorageConfig(t) + basePath := "files/cg_destroy_test/seg1" + + schema := arrow.NewSchema([]arrow.Field{ + { + Name: "pk", + Type: arrow.PrimitiveTypes.Int64, + Nullable: false, + Metadata: arrow.NewMetadata([]string{ArrowFieldIdMetadataKey}, []string{"100"}), + }, + }, nil) + b := array.NewRecordBuilder(memory.DefaultAllocator, schema) + defer b.Release() + b.Field(0).(*array.Int64Builder).Append(int64(0)) + rec := b.NewRecord() + defer rec.Release() + + w, err := NewFFIPackedWriter(basePath, schema, + []storagecommon.ColumnGroup{{Columns: []int{0}, GroupID: storagecommon.DefaultShortColumnGroupID}}, + cfg, nil) + require.NoError(t, err) + require.NoError(t, w.WriteRecordBatch(rec)) + out, err := w.Close() + require.NoError(t, err) + cgs, ok := out.(*ColumnGroups) + require.True(t, ok, "Close should return *ColumnGroups") + require.NotNil(t, cgs.cColumnGroups, "Close should yield non-nil cColumnGroups") + + cgs.Destroy() + require.Nil(t, cgs.cColumnGroups, "first Destroy must clear cColumnGroups") + cgs.Destroy() // second call must be a no-op + assert.Nil(t, cgs.cColumnGroups) + + var nilHandle *ColumnGroups + assert.NotPanics(t, func() { nilHandle.Destroy() }, + "Destroy on nil receiver must be safe") +} + +// TestFFIPackedWriter_DoubleCloseRejected verifies that the second Close +// after a successful first Close returns an error rather than re-running +// the close FFI on an exhausted handle. +func TestFFIPackedWriter_DoubleCloseRejected(t *testing.T) { + cfg := manifestTestStorageConfig(t) + basePath := "files/ffi_double_close/seg1" + + schema := arrow.NewSchema([]arrow.Field{ + { + Name: "pk", + Type: arrow.PrimitiveTypes.Int64, + Nullable: false, + Metadata: arrow.NewMetadata([]string{ArrowFieldIdMetadataKey}, []string{"100"}), + }, + }, nil) + b := array.NewRecordBuilder(memory.DefaultAllocator, schema) + defer b.Release() + b.Field(0).(*array.Int64Builder).Append(int64(0)) + rec := b.NewRecord() + defer rec.Release() + + w, err := NewFFIPackedWriter(basePath, schema, + []storagecommon.ColumnGroup{{Columns: []int{0}, GroupID: storagecommon.DefaultShortColumnGroupID}}, + cfg, nil) + require.NoError(t, err) + require.NoError(t, w.WriteRecordBatch(rec)) + + out, err := w.Close() + require.NoError(t, err) + defer out.Destroy() + + _, err = w.Close() + require.Error(t, err, "second Close must return an error") + require.Contains(t, err.Error(), "already closed") +} diff --git a/internal/storagev2/packed/packed_reader_ffi_test.go b/internal/storagev2/packed/packed_reader_ffi_test.go index eb1e76ba1b..5881f89acb 100644 --- a/internal/storagev2/packed/packed_reader_ffi_test.go +++ b/internal/storagev2/packed/packed_reader_ffi_test.go @@ -93,25 +93,29 @@ func TestPackedFFIReader(t *testing.T) { {Columns: []int{0, 1}, GroupID: storagecommon.DefaultShortColumnGroupID}, } + storageConfig := &indexpb.StorageConfig{ + RootPath: dir, + StorageType: "local", + } + // Create FFI packed writer and write data - pw, err := NewFFIPackedWriter(basePath, version, schema, columnGroups, nil, nil) + pw, err := NewFFIPackedWriter(basePath, schema, columnGroups, storageConfig, nil) require.NoError(t, err) err = pw.WriteRecordBatch(rec) require.NoError(t, err) - manifest, err := pw.Close() + out, err := pw.Close() + require.NoError(t, err) + defer out.Destroy() + + manifest, err := CommitManifestUpdates(basePath, version, storageConfig, + &ManifestUpdates{NewFiles: out}) require.NoError(t, err) require.NotEmpty(t, manifest) t.Logf("Successfully wrote %d rows with %d-dim float vectors, manifest: %s", numRows, dim, manifest) - // Create storage config for reader - storageConfig := &indexpb.StorageConfig{ - RootPath: dir, - StorageType: "local", - } - // Create FFI packed reader neededColumns := []string{"pk", "vector"} reader, err := NewFFIPackedReader(manifest, schema, neededColumns, 8192, storageConfig, nil) @@ -238,22 +242,27 @@ func TestPackedFFIReaderPartialColumns(t *testing.T) { {Columns: []int{0, 1, 2}, GroupID: storagecommon.DefaultShortColumnGroupID}, } - // Write data - pw, err := NewFFIPackedWriter(basePath, version, schema, columnGroups, nil, nil) - require.NoError(t, err) - - err = pw.WriteRecordBatch(rec) - require.NoError(t, err) - - manifest, err := pw.Close() - require.NoError(t, err) - // Create storage config storageConfig := &indexpb.StorageConfig{ RootPath: dir, StorageType: "local", } + // Write data + pw, err := NewFFIPackedWriter(basePath, schema, columnGroups, storageConfig, nil) + require.NoError(t, err) + + err = pw.WriteRecordBatch(rec) + require.NoError(t, err) + + out, err := pw.Close() + require.NoError(t, err) + defer out.Destroy() + + manifest, err := CommitManifestUpdates(basePath, version, storageConfig, + &ManifestUpdates{NewFiles: out}) + require.NoError(t, err) + // Read only pk and score columns (skip vector) neededColumns := []string{"pk", "score"} partialSchema := arrow.NewSchema([]arrow.Field{ @@ -343,6 +352,11 @@ func TestPackedFFIReaderMultipleBatches(t *testing.T) { var manifest string totalWrittenRows := 0 + storageConfig := &indexpb.StorageConfig{ + RootPath: dir, + StorageType: "local", + } + // Write multiple batches for batch := 0; batch < numWrites; batch++ { b := array.NewRecordBuilder(memory.DefaultAllocator, schema) @@ -364,13 +378,18 @@ func TestPackedFFIReaderMultipleBatches(t *testing.T) { rec := b.NewRecord() - pw, err := NewFFIPackedWriter(basePath, version, schema, columnGroups, nil, nil) + pw, err := NewFFIPackedWriter(basePath, schema, columnGroups, storageConfig, nil) require.NoError(t, err) err = pw.WriteRecordBatch(rec) require.NoError(t, err) - manifest, err = pw.Close() + out, err := pw.Close() + require.NoError(t, err) + + manifest, err = CommitManifestUpdates(basePath, version, storageConfig, + &ManifestUpdates{NewFiles: out}) + out.Destroy() require.NoError(t, err) _, version, err = UnmarshalManifestPath(manifest) @@ -382,11 +401,7 @@ func TestPackedFFIReaderMultipleBatches(t *testing.T) { rec.Release() } - // Read all data - storageConfig := &indexpb.StorageConfig{ - RootPath: dir, - StorageType: "local", - } + // Read all data using the same storageConfig neededColumns := []string{"pk", "vector"} reader, err := NewFFIPackedReader(manifest, schema, neededColumns, 8192, storageConfig, nil) diff --git a/internal/storagev2/packed/packed_writer_ffi.go b/internal/storagev2/packed/packed_writer_ffi.go index 9e12231c69..df7cccf153 100644 --- a/internal/storagev2/packed/packed_writer_ffi.go +++ b/internal/storagev2/packed/packed_writer_ffi.go @@ -28,16 +28,15 @@ package packed import "C" import ( + "fmt" "strings" "unsafe" "github.com/apache/arrow/go/v17/arrow" "github.com/apache/arrow/go/v17/arrow/cdata" "github.com/samber/lo" - "go.uber.org/zap" "github.com/milvus-io/milvus/internal/storagecommon" - "github.com/milvus-io/milvus/pkg/v3/log" "github.com/milvus-io/milvus/pkg/v3/proto/indexcgopb" "github.com/milvus-io/milvus/pkg/v3/proto/indexpb" "github.com/milvus-io/milvus/pkg/v3/util/paramtable" @@ -76,7 +75,12 @@ func CreateStorageConfig() *indexpb.StorageConfig { return storageConfig } -func NewFFIPackedWriter(basePath string, baseVersion int64, schema *arrow.Schema, columnGroups []storagecommon.ColumnGroup, storageConfig *indexpb.StorageConfig, storagePluginContext *indexcgopb.StoragePluginContext) (*FFIPackedWriter, error) { +// NewFFIPackedWriter creates a writer that produces parquet files under +// basePath. The writer knows nothing about manifests or versions — its +// only job is to write data files. Close returns the resulting column +// groups, which the caller passes to packed.CommitManifestUpdates to +// register them with a manifest version. +func NewFFIPackedWriter(basePath string, schema *arrow.Schema, columnGroups []storagecommon.ColumnGroup, storageConfig *indexpb.StorageConfig, storagePluginContext *indexcgopb.StoragePluginContext) (*FFIPackedWriter, error) { cBasePath := C.CString(basePath) defer C.free(unsafe.Pointer(cBasePath)) @@ -145,16 +149,16 @@ func NewFFIPackedWriter(basePath string, baseVersion int64, schema *arrow.Schema return &FFIPackedWriter{ basePath: basePath, - baseVersion: baseVersion, cWriterHandle: writerHandle, cProperties: cProperties, }, nil } -// AsNewColumnGroups marks this writer so that Close() calls loon_transaction_add_column_group -// for each written column group instead of loon_transaction_append_files. -// Call this when the written columns are brand-new additions to an existing manifest -// (e.g. function-field backfill), not when appending more files to the same column structure. +// AsNewColumnGroups marks this writer so that the column groups returned +// by Close should be staged via loon_transaction_add_column_group instead +// of loon_transaction_append_files when later passed to +// CommitManifestUpdates. Use true when adding columns that do not yet +// exist in the manifest (e.g. function-field backfill). func (pw *FFIPackedWriter) AsNewColumnGroups() *FFIPackedWriter { pw.addNewColumnGroups = true return pw @@ -176,52 +180,78 @@ func (pw *FFIPackedWriter) WriteRecordBatch(recordBatch arrow.Record) error { return HandleLoonFFIResult(result) } -func (pw *FFIPackedWriter) Close() (string, error) { - var cColumnGroups *C.LoonColumnGroups +// ColumnGroups is the data carrier returned by FFIPackedWriter.Close. It +// holds the column-groups payload produced by the C writer and owns C +// memory; the caller MUST call Destroy after passing the handle to +// CommitManifestUpdates (success or failure). Destroy is idempotent; +// a nil cColumnGroups indicates the handle has already been released. +type ColumnGroups struct { + cColumnGroups *C.LoonColumnGroups + addNewColumnGroups bool +} - result := C.loon_writer_close(pw.cWriterHandle, nil, nil, 0, &cColumnGroups) - if err := HandleLoonFFIResult(result); err != nil { - return "", err +// Destroy releases C memory. Safe to call multiple times. +func (f *ColumnGroups) Destroy() { + if f == nil || f.cColumnGroups == nil { + return } + C.loon_column_groups_destroy(f.cColumnGroups) + f.cColumnGroups = nil +} - cBasePath := C.CString(pw.basePath) - defer C.free(unsafe.Pointer(cBasePath)) - var transationHandle C.LoonTransactionHandle - - result = C.loon_transaction_begin(cBasePath, pw.cProperties, C.int64_t(pw.baseVersion), C.LOON_TRANSACTION_RESOLVE_OVERWRITE /* resolve_id */, getRetryLimit() /* retry_limit */, &transationHandle) - if err := HandleLoonFFIResult(result); err != nil { - return "", err +// applyTo stages the column groups onto a loon transaction. +// +// When addNewColumnGroups is true the groups are added one-by-one via +// loon_transaction_add_column_group (function-backfill case where the +// schema is being extended). Otherwise they are appended in one call via +// loon_transaction_append_files (normal multi-batch write case). +func (f *ColumnGroups) applyTo(handle C.LoonTransactionHandle) error { + if f == nil || f.cColumnGroups == nil { + return nil } - defer C.loon_transaction_destroy(transationHandle) - - if pw.addNewColumnGroups { - // Each written group is a brand-new column group: add them one-by-one. - // loon_transaction_append_files requires the count to match existing groups, - // which would fail when extending the schema (e.g. backfill adds sparse vector - // to a manifest that already has 2 groups of different columns). - numGroups := int(cColumnGroups.num_of_column_groups) - colGroupSlice := (*[1 << 20]C.LoonColumnGroup)(unsafe.Pointer(cColumnGroups.column_group_array))[:numGroups:numGroups] - for i := range colGroupSlice { - result = C.loon_transaction_add_column_group(transationHandle, &colGroupSlice[i]) - if err := HandleLoonFFIResult(result); err != nil { - return "", err + if f.addNewColumnGroups { + num := int(f.cColumnGroups.num_of_column_groups) + slice := unsafe.Slice(f.cColumnGroups.column_group_array, num) + for i := range slice { + if err := HandleLoonFFIResult(C.loon_transaction_add_column_group(handle, &slice[i])); err != nil { + return fmt.Errorf("commit manifest add_column_group: %w", err) } } - } else { - result = C.loon_transaction_append_files(transationHandle, cColumnGroups) - if err := HandleLoonFFIResult(result); err != nil { - return "", err - } + return nil } - - var cCommitVersion C.int64_t - result = C.loon_transaction_commit(transationHandle, &cCommitVersion) - if err := HandleLoonFFIResult(result); err != nil { - return "", err + if err := HandleLoonFFIResult(C.loon_transaction_append_files(handle, f.cColumnGroups)); err != nil { + return fmt.Errorf("commit manifest append_files: %w", err) } + return nil +} - log.Info("FFI writer closed", zap.Int64("version", int64(cCommitVersion))) - - defer C.loon_properties_free(pw.cProperties) - return MarshalManifestPath(pw.basePath, int64(cCommitVersion)), nil +// Close closes the underlying loon writer and returns the column-groups +// payload. The writer never touches the manifest — the caller is +// responsible for passing the returned handle to CommitManifestUpdates +// and calling Destroy when done. +// +// Close releases the writer's C resources (loon writer handle and +// cProperties) in a defer, so even when loon_writer_close fails those +// resources are reclaimed. After Close the writer is exhausted; further +// Close or Write calls fail. +func (pw *FFIPackedWriter) Close() (WriterOutput, error) { + if pw.closed { + return nil, fmt.Errorf("FFIPackedWriter already closed") + } + pw.closed = true + defer func() { + if pw.cProperties != nil { + C.loon_properties_free(pw.cProperties) + pw.cProperties = nil + } + }() + var cColumnGroups *C.LoonColumnGroups + result := C.loon_writer_close(pw.cWriterHandle, nil, nil, 0, &cColumnGroups) + if err := HandleLoonFFIResult(result); err != nil { + return nil, err + } + return &ColumnGroups{ + cColumnGroups: cColumnGroups, + addNewColumnGroups: pw.addNewColumnGroups, + }, nil } diff --git a/internal/storagev2/packed/packed_writer_ffi_test.go b/internal/storagev2/packed/packed_writer_ffi_test.go index 924d9fb07e..c0e014527c 100644 --- a/internal/storagev2/packed/packed_writer_ffi_test.go +++ b/internal/storagev2/packed/packed_writer_ffi_test.go @@ -93,15 +93,21 @@ func TestPackedFFIWriter(t *testing.T) { } // Create FFI packed writer - pw, err := NewFFIPackedWriter(basePath, version, schema, columnGroups, nil, nil) + cfg := CreateStorageConfig() + pw, err := NewFFIPackedWriter(basePath, schema, columnGroups, cfg, nil) require.NoError(t, err) // Write record batch err = pw.WriteRecordBatch(rec) require.NoError(t, err) - // Close writer and get manifest - manifest, err := pw.Close() + // Close writer to obtain column groups, commit via manifest update. + out, err := pw.Close() + require.NoError(t, err) + + manifest, err := CommitManifestUpdates(basePath, version, cfg, + &ManifestUpdates{NewFiles: out}) + out.Destroy() require.NoError(t, err) require.NotEmpty(t, manifest) @@ -114,3 +120,56 @@ func TestPackedFFIWriter(t *testing.T) { t.Logf("Successfully wrote %d rows with %d-dim float vectors, manifest: %s", numRows, dim, manifest) } } + +func TestFFIPackedWriter_CloseThenCommitUpdates(t *testing.T) { + paramtable.Init() + pt := paramtable.Get() + pt.Save(pt.CommonCfg.StorageType.Key, "local") + dir := t.TempDir() + pt.Save(pt.LocalStorageCfg.Path.Key, dir) + t.Cleanup(func() { + pt.Reset(pt.CommonCfg.StorageType.Key) + pt.Reset(pt.LocalStorageCfg.Path.Key) + }) + + schema := arrow.NewSchema([]arrow.Field{ + { + Name: "pk", + Type: arrow.PrimitiveTypes.Int64, + Nullable: false, + Metadata: arrow.NewMetadata([]string{ArrowFieldIdMetadataKey}, []string{"100"}), + }, + }, nil) + + b := array.NewRecordBuilder(memory.DefaultAllocator, schema) + defer b.Release() + pkb := b.Field(0).(*array.Int64Builder) + for i := 0; i < 4; i++ { + pkb.Append(int64(i)) + } + rec := b.NewRecord() + defer rec.Release() + + columnGroups := []storagecommon.ColumnGroup{ + {Columns: []int{0}, GroupID: storagecommon.DefaultShortColumnGroupID}, + } + + basePath := "files/close_commit_test/1" + cfg := CreateStorageConfig() + w, err := NewFFIPackedWriter(basePath, schema, columnGroups, cfg, nil) + require.NoError(t, err) + require.NoError(t, w.WriteRecordBatch(rec)) + + out, err := w.Close() + require.NoError(t, err) + require.NotNil(t, out) + defer out.Destroy() + + mfPath, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg, + &ManifestUpdates{NewFiles: out}) + require.NoError(t, err) + + _, v, err := UnmarshalManifestPath(mfPath) + require.NoError(t, err) + require.Equal(t, int64(1), v, "exactly one version bump expected") +} diff --git a/internal/storagev2/packed/segment_writer_ffi.go b/internal/storagev2/packed/segment_writer_ffi.go index c6d41963a2..f4307c1bd3 100644 --- a/internal/storagev2/packed/segment_writer_ffi.go +++ b/internal/storagev2/packed/segment_writer_ffi.go @@ -44,29 +44,20 @@ type TextColumnConfig struct { RewriteMode bool // true = input is LOB references, decode & rewrite during compaction } -// SegmentWriterConfig represents configuration for SegmentWriter. -// ReadVersion and RetryLimit are used by the Go-layer transaction logic, +// SegmentWriterConfig represents configuration for SegmentWriter. The +// writer is concerned only with file output; manifest-level concerns +// (version, retry) live in CommitManifestUpdates. type SegmentWriterConfig struct { SegmentPath string - ReadVersion int64 // manifest version for transaction (Go-layer only) - RetryLimit uint32 // transaction retry limit (Go-layer only) TextColumns []TextColumnConfig } -// SegmentWriterResult contains the result of closing a SegmentWriter. -type SegmentWriterResult struct { - ManifestPath string - CommittedVersion int64 - RowsWritten int64 -} - // FFISegmentWriter wraps the C SegmentWriter handle for incremental writes. type FFISegmentWriter struct { handle C.LoonSegmentWriterHandle cProperties *C.LoonProperties schema *arrow.Schema - basePath string // segment base path for transaction - readVersion int64 // manifest read version for transaction + closed bool } // NewFFISegmentWriter creates a new segment writer via FFI. @@ -107,8 +98,6 @@ func NewFFISegmentWriter( handle: writerHandle, cProperties: cProperties, schema: schema, - basePath: config.SegmentPath, - readVersion: config.ReadVersion, }, nil } @@ -127,85 +116,97 @@ func (w *FFISegmentWriter) Write(record arrow.Record) error { return HandleLoonFFIResult(result) } -// Flush flushes buffered data to storage. -func (w *FFISegmentWriter) Flush() error { +// SyncBuffered syncs buffered data to storage without closing the writer. +// Used for mid-stream flushes; Close detaches output and finalizes the writer. +func (w *FFISegmentWriter) SyncBuffered() error { result := C.loon_segment_writer_flush(w.handle) return HandleLoonFFIResult(result) } -// Close closes the writer and commits the manifest via Transaction. -// Returns SegmentWriterResult with the committed manifest path. -// Pattern matches FFIPackedWriter.Close(): C++ writer returns ColumnGroups + LobFiles, -// Go layer handles Transaction begin → append_files → add_lob_files → commit. -func (w *FFISegmentWriter) Close() (*SegmentWriterResult, error) { - var cOutput C.LoonSegmentWriteOutput +// SegmentOutput is the data carrier returned by FFISegmentWriter.Close. +// It holds the column-groups + LOB payload produced by the C writer and +// owns C memory; the caller MUST call Destroy after passing the handle to +// CommitManifestUpdates (success or failure). Destroy is idempotent; a +// nil column_groups pointer indicates the handle has already been released. +type SegmentOutput struct { + cOutput C.LoonSegmentWriteOutput + rowsWritten int64 +} +// RowsWritten returns the number of rows the segment writer reported. +func (f *SegmentOutput) RowsWritten() int64 { return f.rowsWritten } + +// Destroy releases the C output (LOB file strings + array). Safe to call +// multiple times — the nil column_groups pointer left behind serves as +// the already-released marker. +func (f *SegmentOutput) Destroy() { + if f == nil || f.cOutput.column_groups == nil { + return + } + C.loon_segment_write_output_free(&f.cOutput) + f.cOutput.column_groups = nil + f.cOutput.lob_files = nil + f.cOutput.num_lob_files = 0 +} + +// applyTo stages the column-groups + LOB payload onto a loon transaction. +// Column groups (when present) are appended via loon_transaction_append_files, +// and each LOB file is registered via loon_transaction_add_lob_file. +func (f *SegmentOutput) applyTo(handle C.LoonTransactionHandle) error { + if f == nil { + return nil + } + if f.cOutput.column_groups != nil { + if err := HandleLoonFFIResult(C.loon_transaction_append_files(handle, f.cOutput.column_groups)); err != nil { + return fmt.Errorf("commit manifest append_files (segment): %w", err) + } + } + if f.cOutput.num_lob_files > 0 && f.cOutput.lob_files != nil { + lob := unsafe.Slice(f.cOutput.lob_files, f.cOutput.num_lob_files) + for i := range lob { + if err := HandleLoonFFIResult(C.loon_transaction_add_lob_file(handle, &lob[i])); err != nil { + return fmt.Errorf("commit manifest add_lob_file: %w", err) + } + } + } + return nil +} + +// Close closes the underlying segment writer and returns the column-groups +// + LOB payload. The writer never touches the manifest — the caller is +// responsible for passing the returned handle to CommitManifestUpdates +// and calling Destroy on the returned WriterOutput when done. +// +// Close releases the writer's C resources (segment-writer handle and +// cProperties) in a defer, so even when loon_segment_writer_close fails +// those resources are reclaimed. After Close the writer is exhausted; +// further Close or Write calls fail. +func (w *FFISegmentWriter) Close() (WriterOutput, error) { + if w.closed { + return nil, fmt.Errorf("FFISegmentWriter already closed") + } + w.closed = true + defer func() { + if w.handle != 0 { + C.loon_segment_writer_destroy(w.handle) + w.handle = 0 + } + if w.cProperties != nil { + C.loon_properties_free(w.cProperties) + w.cProperties = nil + } + }() + var cOutput C.LoonSegmentWriteOutput result := C.loon_segment_writer_close(w.handle, &cOutput) if err := HandleLoonFFIResult(result); err != nil { return nil, err } - - rowsWritten := int64(cOutput.rows_written) - - cBasePath := C.CString(w.basePath) - defer C.free(unsafe.Pointer(cBasePath)) - - var transactionHandle C.LoonTransactionHandle - result = C.loon_transaction_begin(cBasePath, w.cProperties, - C.int64_t(w.readVersion), C.int32_t(0), getRetryLimit(), &transactionHandle) - if err := HandleLoonFFIResult(result); err != nil { - return nil, err - } - defer C.loon_transaction_destroy(transactionHandle) - - // append column groups - if cOutput.column_groups != nil { - result = C.loon_transaction_append_files(transactionHandle, cOutput.column_groups) - if err := HandleLoonFFIResult(result); err != nil { - return nil, err - } - } - - // add LOB files - if cOutput.num_lob_files > 0 && cOutput.lob_files != nil { - cLobSlice := unsafe.Slice(cOutput.lob_files, cOutput.num_lob_files) - for _, cLob := range cLobSlice { - result = C.loon_transaction_add_lob_file(transactionHandle, &cLob) - if err := HandleLoonFFIResult(result); err != nil { - return nil, err - } - } - } - - // commit - var cCommitVersion C.int64_t - result = C.loon_transaction_commit(transactionHandle, &cCommitVersion) - if err := HandleLoonFFIResult(result); err != nil { - return nil, err - } - - // free C output (LOB file strings + array) - C.loon_segment_write_output_free(&cOutput) - - return &SegmentWriterResult{ - ManifestPath: MarshalManifestPath(w.basePath, int64(cCommitVersion)), - CommittedVersion: int64(cCommitVersion), - RowsWritten: rowsWritten, + return &SegmentOutput{ + cOutput: cOutput, + rowsWritten: int64(cOutput.rows_written), }, nil } -// Destroy destroys the writer and releases resources. -func (w *FFISegmentWriter) Destroy() { - if w.handle != 0 { - C.loon_segment_writer_destroy(w.handle) - w.handle = 0 - } - if w.cProperties != nil { - C.loon_properties_free(w.cProperties) - w.cProperties = nil - } -} - // buildCSegmentWriterConfig converts Go config to C config. func buildCSegmentWriterConfig(config *SegmentWriterConfig) *C.LoonSegmentWriterConfig { cConfig := (*C.LoonSegmentWriterConfig)(C.malloc(C.sizeof_LoonSegmentWriterConfig)) diff --git a/internal/storagev2/packed/transaction_test.go b/internal/storagev2/packed/transaction_test.go index ed53658ef8..424d106bc3 100644 --- a/internal/storagev2/packed/transaction_test.go +++ b/internal/storagev2/packed/transaction_test.go @@ -62,7 +62,7 @@ func createBaseManifest(t *testing.T, basePath string, storageConfig *indexpb.St {Columns: []int{0, 1}, GroupID: storagecommon.DefaultShortColumnGroupID}, } - pw, err := NewFFIPackedWriter(basePath, 0, schema, columnGroups, storageConfig, nil) + pw, err := NewFFIPackedWriter(basePath, schema, columnGroups, storageConfig, nil) require.NoError(t, err) b := array.NewRecordBuilder(memory.DefaultAllocator, schema) @@ -75,7 +75,12 @@ func createBaseManifest(t *testing.T, basePath string, storageConfig *indexpb.St err = pw.WriteRecordBatch(rec) require.NoError(t, err) - manifestPath, err := pw.Close() + out, err := pw.Close() + require.NoError(t, err) + defer out.Destroy() + + manifestPath, err := CommitManifestUpdates(basePath, ManifestEarliest, storageConfig, + &ManifestUpdates{NewFiles: out}) require.NoError(t, err) return manifestPath } diff --git a/internal/storagev2/packed/type.go b/internal/storagev2/packed/type.go index 7c9214f6e2..47294fd564 100644 --- a/internal/storagev2/packed/type.go +++ b/internal/storagev2/packed/type.go @@ -37,15 +37,16 @@ type PackedWriter struct { type FFIPackedWriter struct { basePath string - baseVersion int64 cWriterHandle C.LoonWriterHandle cProperties *C.LoonProperties - // addNewColumnGroups controls whether Close() uses loon_transaction_add_column_group - // (true) or loon_transaction_append_files (false). - // - // Use true when adding columns that do not yet exist in the manifest (e.g. backfill). - // Use false (default) when writing more files into the same column structure (e.g. multi-batch write). + // addNewColumnGroups is propagated into the ColumnGroups output of + // Close so CommitManifestUpdates can choose between + // loon_transaction_add_column_group (true) and + // loon_transaction_append_files (false). Use true when adding columns + // that do not yet exist in the manifest (e.g. function backfill). addNewColumnGroups bool + // closed is set once Close has consumed the underlying C writer handle. + closed bool } type PackedReader struct {