chore: import upstream snapshot with attribution

internal/streamingnode/server/flusher/flusherimpl/msg_handler_impl.go

@@ -0,0 +1,171 @@
+// 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
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package flusherimpl
+
+import (
+	"context"
+
+	"github.com/cockroachdb/errors"
+
+	"github.com/milvus-io/milvus/internal/flushcommon/writebuffer"
+	"github.com/milvus-io/milvus/internal/streamingnode/server/resource"
+	"github.com/milvus-io/milvus/pkg/v3/mlog"
+	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
+	"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/v3/util/merr"
+	"github.com/milvus-io/milvus/pkg/v3/util/retry"
+)
+
+func newMsgHandler(wbMgr writebuffer.BufferManager) *msgHandlerImpl {
+	return &msgHandlerImpl{
+		wbMgr: wbMgr,
+	}
+}
+
+type msgHandlerImpl struct {
+	wbMgr writebuffer.BufferManager
+}
+
+func (impl *msgHandlerImpl) HandleCreateSegment(ctx context.Context, createSegmentMsg message.ImmutableCreateSegmentMessageV2) error {
+	vchannel := createSegmentMsg.VChannel()
+	h := createSegmentMsg.Header()
+
+	if err := impl.createNewGrowingSegment(ctx, vchannel, h); err != nil {
+		return err
+	}
+	logger := mlog.With(mlog.FieldMessage(createSegmentMsg))
+
+	if err := impl.wbMgr.CreateNewGrowingSegment(ctx, vchannel, h.PartitionId, h.SegmentId, h.SchemaVersion); err != nil {
+		logger.Warn(ctx, "fail to create new growing segment")
+		return err
+	}
+	mlog.Info(ctx, "create new growing segment")
+	return nil
+}
+
+func (impl *msgHandlerImpl) createNewGrowingSegment(ctx context.Context, vchannel string, h *message.CreateSegmentMessageHeader) error {
+	if h.Level == datapb.SegmentLevel_L0 {
+		// L0 segment should not be flushed directly, but not create than flush.
+		// the create segment operation is used to protect the binlog from garbage collection.
+		// L0 segment's binlog upload and flush operation is handled once.
+		// so we can skip the create segment operation here. (not strict promise exactly)
+		return nil
+	}
+	// Transfer the pending segment into growing state.
+	// Alloc the growing segment at datacoord first.
+	mix, err := resource.Resource().MixCoordClient().GetWithContext(ctx)
+	if err != nil {
+		return err
+	}
+	logger := mlog.With(mlog.FieldCollectionID(h.CollectionId), mlog.FieldPartitionID(h.PartitionId), mlog.FieldSegmentID(h.SegmentId))
+	return retry.Do(ctx, func() (err error) {
+		// TODO: propagate SchemaVersion from CreateSegmentMessageHeader into AllocSegmentRequest
+		// so that DataCoord records the correct schema version for streaming-created segments.
+		// Without this, new segments get SchemaVersion=0 and will be falsely flagged for backfill.
+		// Tracked in companion PR: https://github.com/milvus-io/milvus/pull/48865
+		resp, err := mix.AllocSegment(ctx, &datapb.AllocSegmentRequest{
+			CollectionId:         h.CollectionId,
+			PartitionId:          h.PartitionId,
+			SegmentId:            h.SegmentId,
+			Vchannel:             vchannel,
+			StorageVersion:       h.StorageVersion,
+			IsCreatedByStreaming: true,
+			SchemaVersion:        h.SchemaVersion,
+		})
+		if err := merr.CheckRPCCall(resp, err); err != nil {
+			logger.Warn(ctx, "failed to alloc growing segment at datacoord")
+			return errors.Wrap(err, "failed to alloc growing segment at datacoord")
+		}
+		logger.Info(ctx, "alloc growing segment at datacoord", mlog.Int32("schemaVersion", h.SchemaVersion))
+		return nil
+	}, retry.AttemptAlways(), retry.RetryErr(func(error) bool { return true }))
+}
+
+func (impl *msgHandlerImpl) HandleFlush(flushMsg message.ImmutableFlushMessageV2) error {
+	vchannel := flushMsg.VChannel()
+	if err := impl.wbMgr.SealSegments(context.Background(), vchannel, []int64{flushMsg.Header().SegmentId}); err != nil {
+		return errors.Wrap(err, "failed to seal segments")
+	}
+	return nil
+}
+
+func (impl *msgHandlerImpl) HandleManualFlush(flushMsg message.ImmutableManualFlushMessageV2) error {
+	vchannel := flushMsg.VChannel()
+	if err := impl.wbMgr.SealSegments(context.Background(), vchannel, flushMsg.Header().SegmentIds); err != nil {
+		return errors.Wrap(err, "failed to seal segments")
+	}
+	if err := impl.wbMgr.FlushChannel(context.Background(), vchannel, flushMsg.TimeTick()); err != nil {
+		return errors.Wrap(err, "failed to flush channel")
+	} // may be redundant.
+	return nil
+}
+
+func (impl *msgHandlerImpl) HandleFlushAll(vchannel string, flushAllMsg message.ImmutableFlushAllMessageV2) error {
+	if err := impl.wbMgr.SealAllSegments(context.Background(), vchannel); err != nil {
+		return errors.Wrap(err, "failed to seal all segments")
+	}
+	// Use FlushAllMsg's ts as flush ts.
+	if err := impl.wbMgr.FlushChannel(context.Background(), vchannel, flushAllMsg.TimeTick()); err != nil {
+		return errors.Wrap(err, "failed to flush channel")
+	} // may be redundant.
+	return nil
+}
+
+func (impl *msgHandlerImpl) HandleSchemaChange(ctx context.Context, msg message.ImmutableSchemaChangeMessageV2) error {
+	return impl.wbMgr.SealSegments(context.Background(), msg.VChannel(), msg.Header().FlushedSegmentIds)
+}
+
+func (impl *msgHandlerImpl) HandleAlterCollection(ctx context.Context, putCollectionMsg message.ImmutableAlterCollectionMessageV2) error {
+	return impl.wbMgr.SealSegments(context.Background(), putCollectionMsg.VChannel(), putCollectionMsg.Header().FlushedSegmentIds)
+}
+
+func (impl *msgHandlerImpl) HandleTruncateCollection(flushMsg message.ImmutableTruncateCollectionMessageV2) error {
+	vchannel := flushMsg.VChannel()
+	if err := impl.wbMgr.SealSegments(context.Background(), vchannel, flushMsg.Header().SegmentIds); err != nil {
+		return errors.Wrap(err, "failed to seal segments")
+	}
+	if err := impl.wbMgr.FlushChannel(context.Background(), vchannel, flushMsg.TimeTick()); err != nil {
+		return errors.Wrap(err, "failed to flush channel")
+	}
+	return nil
+}
+
+// HandleAlterWAL handles Alter WAL message by sealing all segments and flushing the channel.
+// This ensures all buffered data is persisted before switching to the new WAL implementation.
+func (impl *msgHandlerImpl) HandleAlterWAL(ctx context.Context, alterWALMsg message.ImmutableAlterWALMessageV2, currentVChannel string) error {
+	vchannel := currentVChannel
+	logger := mlog.With(
+		mlog.FieldVChannel(vchannel),
+		mlog.Uint64("alterWALTimeTick", alterWALMsg.TimeTick()),
+		mlog.String("currentWALName", alterWALMsg.WALName().String()),
+		mlog.Stringer("targetWALName", alterWALMsg.Header().TargetWalName))
+
+	// Seal all segments in the current vchannel before WAL switch
+	if err := impl.wbMgr.SealAllSegments(ctx, vchannel); err != nil {
+		logger.Warn(ctx, "failed to seal all segments for WAL switch", mlog.Err(err))
+		return errors.Wrap(err, "failed to seal all segments")
+	}
+	logger.Info(ctx, "sealed all segments for WAL switch")
+
+	// Flush channel to persist buffered data before switching WAL
+	if err := impl.wbMgr.FlushChannel(ctx, vchannel, alterWALMsg.TimeTick()); err != nil {
+		logger.Warn(ctx, "failed to flush channel for WAL switch", mlog.Err(err))
+		return errors.Wrap(err, "failed to flush channel")
+	}
+	logger.Info(ctx, "flushed channel for WAL switch")
+	return nil
+}