milvus/internal/streamingnode/server/wal/adaptor/wal_adaptor.go

package adaptor

import (
	"context"
	"time"

	"github.com/cenkalti/backoff/v4"
	"github.com/cockroachdb/errors"
	"go.opentelemetry.io/otel/codes"
	"go.uber.org/atomic"
	"google.golang.org/protobuf/types/known/anypb"

	"github.com/milvus-io/milvus/internal/streamingnode/server/flusher/flusherimpl"
	"github.com/milvus-io/milvus/internal/streamingnode/server/resource"
	"github.com/milvus-io/milvus/internal/streamingnode/server/wal"
	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/adaptor/rate"
	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors"
	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/metricsutil"
	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/utility"
	"github.com/milvus-io/milvus/internal/util/streamingutil/status"
	"github.com/milvus-io/milvus/pkg/v3/mlog"
	"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
	"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
	"github.com/milvus-io/milvus/pkg/v3/streaming/walimpls"
	"github.com/milvus-io/milvus/pkg/v3/util/conc"
	"github.com/milvus-io/milvus/pkg/v3/util/contextutil"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

var _ wal.WAL = (*walAdaptorImpl)(nil)

type gracefulCloseFunc func()

// adaptImplsToROWAL creates a new readonly wal from wal impls.
func adaptImplsToROWAL(
	basicWAL walimpls.WALImpls,
	cleanup func(),
) *roWALAdaptorImpl {
	logger := resource.Resource().Logger().With(
		mlog.FieldComponent("wal"),
		mlog.String("channel", basicWAL.Channel().String()),
	)
	ctx, cancel := context.WithCancel(context.Background()) //nolint:gosec // cancel is stored in availableCancel and called in Close()
	roWAL := &roWALAdaptorImpl{
		WALRateLimitComponent: rate.NewWALRateLimitComponent(basicWAL.Channel()),

		roWALImpls:      basicWAL,
		lifetime:        typeutil.NewLifetime(),
		availableCtx:    ctx,
		availableCancel: cancel,
		idAllocator:     typeutil.NewIDAllocator(),
		scannerRegistry: scannerRegistry{
			channel:     basicWAL.Channel(),
			idAllocator: typeutil.NewIDAllocator(),
		},
		scanners:    typeutil.NewConcurrentMap[int64, wal.Scanner](),
		cleanup:     cleanup,
		scanMetrics: metricsutil.NewScanMetrics(basicWAL.Channel()),
	}
	roWAL.SetLogger(logger)
	return roWAL
}

// adaptImplsToRWWAL creates a new wal from wal impls.
func adaptImplsToRWWAL(
	roWAL *roWALAdaptorImpl,
	builders []interceptors.InterceptorBuilder,
	interceptorParam *interceptors.InterceptorBuildParam,
	flusher *flusherimpl.WALFlusherImpl,
) *walAdaptorImpl {
	if roWAL.Channel().AccessMode != types.AccessModeRW {
		panic("wal should be read-write")
	}
	// build append interceptor for a wal.
	wal := &walAdaptorImpl{
		roWALAdaptorImpl: roWAL,
		rwWALImpls:       roWAL.roWALImpls.(walimpls.WALImpls),
		// TODO: remove the pool, use a queue instead.
		appendExecutionPool:    conc.NewPool[struct{}](0),
		param:                  interceptorParam,
		interceptorBuildResult: buildInterceptor(builders, interceptorParam),
		flusher:                flusher,
		writeMetrics:           metricsutil.NewWriteMetrics(roWAL.Channel(), roWAL.WALName()),
		isFenced:               atomic.NewBool(false),
		appendRateCounter:      utility.NewAverageRateCounter(10 * time.Second), // 10 second sliding window
	}
	wal.writeMetrics.SetLogger(wal.Logger())
	interceptorParam.WAL.Set(wal)
	wal.RegisterMemoryObserver()
	wal.RegisterAppendRateObserver(wal.appendRateCounter)
	return wal
}

// walAdaptorImpl is a wrapper of WALImpls to extend it into a WAL interface.
type walAdaptorImpl struct {
	*roWALAdaptorImpl

	rwWALImpls             walimpls.WALImpls
	appendExecutionPool    *conc.Pool[struct{}]
	param                  *interceptors.InterceptorBuildParam
	interceptorBuildResult interceptorBuildResult
	flusher                *flusherimpl.WALFlusherImpl
	writeMetrics           *metricsutil.WriteMetrics
	isFenced               *atomic.Bool
	appendRateCounter      *utility.AverageRateCounter // tracks append rate (bytes/sec)
}

// Metrics returns the metrics of the wal.
func (w *walAdaptorImpl) Metrics() types.WALMetrics {
	currentMVCC := w.param.MVCCManager.GetMVCCOfVChannel(w.Channel().Name)
	recoveryMetrics := w.flusher.Metrics()
	return types.RWWALMetrics{
		ChannelInfo:      w.Channel(),
		MVCCTimeTick:     currentMVCC.Timetick,
		RecoveryTimeTick: recoveryMetrics.RecoveryTimeTick,
	}
}

// GetLatestMVCCTimestamp get the latest mvcc timestamp of the wal at vchannel.
func (w *walAdaptorImpl) GetLatestMVCCTimestamp(ctx context.Context, vchannel string) (uint64, error) {
	if !w.lifetime.Add(typeutil.LifetimeStateWorking) {
		return 0, status.NewOnShutdownError("wal is on shutdown")
	}
	defer w.lifetime.Done()
	currentMVCC := w.param.MVCCManager.GetMVCCOfVChannel(vchannel)
	if !currentMVCC.Confirmed {
		// if the mvcc is not confirmed, trigger a sync operation to make it confirmed as soon as possible.
		resource.Resource().TimeTickInspector().TriggerSync(w.rwWALImpls.Channel(), false)
	}
	return currentMVCC.Timetick, nil
}

// GetReplicateCheckpoint returns the replicate checkpoint of the wal.
func (w *walAdaptorImpl) GetReplicateCheckpoint() (*utility.ReplicateCheckpoint, error) {
	if !w.lifetime.Add(typeutil.LifetimeStateWorking) {
		return nil, status.NewOnShutdownError("wal is on shutdown")
	}
	defer w.lifetime.Done()

	return w.param.ReplicateManager.GetReplicateCheckpoint()
}

// GetSalvageCheckpoint returns all salvage checkpoints captured during force promote.
func (w *walAdaptorImpl) GetSalvageCheckpoint() []*utility.ReplicateCheckpoint {
	if !w.lifetime.Add(typeutil.LifetimeStateWorking) {
		return nil
	}
	defer w.lifetime.Done()

	return w.param.ReplicateManager.GetSalvageCheckpoint()
}

// Append writes a record to the log.
func (w *walAdaptorImpl) Append(ctx context.Context, msg message.MutableMessage) (_ *wal.AppendResult, err error) {
	if !w.lifetime.Add(typeutil.LifetimeStateWorking) {
		return nil, status.NewOnShutdownError("wal is on shutdown")
	}
	defer w.lifetime.Done()

	ctx, span := message.StartSpanForMessage(ctx, msg, message.SpanNameWALAppend)
	defer func() {
		if err != nil {
			span.RecordError(err)
			span.SetStatus(codes.Error, err.Error())
		}
		span.End()
	}()

	if w.isFenced.Load() {
		// if the wal is fenced, we should reject all append operations.
		return nil, status.NewChannelFenced(w.Channel().String())
	}

	if msg.MessageType().IsDMLMessageType() && w.IsRejected() {
		// if the wal is rate limit rejected, we reject all the DML operation to protect the wal from being overloaded.
		return nil, status.NewRateLimitRejected("")
	}

	// Check if interceptor is ready.
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-w.availableCtx.Done():
		return nil, status.NewOnShutdownError("wal is on shutdown")
	case <-w.interceptorBuildResult.Interceptor.Ready():
	}

	// Setup the term of wal.
	msg = msg.WithWALTerm(w.Channel().Term)

	// we need to promise the state of wal kept consistent with the memory state of streamingnode.
	// So we don't allow the append operation can be canceled by the append caller to avoid leave a inconsistent state of alive wal,
	// the wal append operation can only be canceled when the wal is shutting down.
	ctx, cancel := contextutil.MergeContext(context.WithoutCancel(ctx), w.availableCtx)
	defer cancel()

	appendMetrics := w.writeMetrics.StartAppend(msg)
	ctx = utility.WithAppendMetricsContext(ctx, appendMetrics)

	// Metrics for append message.
	metricsGuard := appendMetrics.StartAppendGuard()

	// Execute the interceptor and wal append.
	var extraAppendResult utility.ExtraAppendResult
	ctx = utility.WithExtraAppendResult(ctx, &extraAppendResult)
	messageID, err := w.interceptorBuildResult.Interceptor.DoAppend(ctx, msg,
		func(ctx context.Context, msg message.MutableMessage) (message.MessageID, error) {
			if notPersistHint := utility.GetNotPersisted(ctx); notPersistHint != nil {
				// do not persist the message if the hint is set.
				return notPersistHint.MessageID, nil
			}
			metricsGuard.StartWALImplAppend()
			msgID, err := w.retryAppendWhenRecoverableError(ctx, msg)
			metricsGuard.FinishWALImplAppend()
			return msgID, err
		})
	metricsGuard.FinishAppend()
	if err != nil {
		appendMetrics.Done(ctx, nil, err)
		if errors.Is(err, walimpls.ErrFenced) {
			// if the append operation of wal is fenced, we should report the error to the client.
			if w.isFenced.CompareAndSwap(false, true) {
				w.forceCancelAfterGracefulTimeout()
				w.Logger().Warn(ctx, "wal is fenced, mark as unavailable, all append opertions will be rejected", mlog.Err(err))
			}
			return nil, status.NewChannelFenced(w.Channel().String())
		}
		return nil, err
	}
	// Mark WAL as fenced if alter WAL message is appended successfully
	// This prevents further append operations during WAL switch
	if msg.MessageType() == message.MessageTypeAlterWAL {
		w.Logger().Info(ctx, "alter WAL message appended, marking WAL as fenced")
		w.isFenced.CompareAndSwap(false, true)
		w.forceCancelAfterGracefulTimeout()
		w.Logger().Info(ctx, "WAL marked as fenced for WAL switch, all append operations will be rejected")
	}
	w.appendRateCounter.Add(int64(msg.EstimateSize()))

	var extra *anypb.Any
	if extraAppendResult.Extra != nil {
		var err error
		if extra, err = anypb.New(extraAppendResult.Extra); err != nil {
			panic("unreachable: failed to marshal extra append result")
		}
	}

	// unwrap the messageID if needed.
	r := &wal.AppendResult{
		MessageID:              messageID,
		LastConfirmedMessageID: extraAppendResult.LastConfirmedMessageID,
		TimeTick:               extraAppendResult.TimeTick,
		TxnCtx:                 extraAppendResult.TxnCtx,
		Extra:                  extra,
	}
	appendMetrics.Done(ctx, r, nil)
	return r, nil
}

// Read overrides the roWALAdaptorImpl.Read to automatically add the append rate counter.
func (w *walAdaptorImpl) Read(ctx context.Context, opts wal.ReadOption) (wal.Scanner, error) {
	// Automatically add the append rate counter to the read options.
	opts.AppendRateCounter = w.appendRateCounter
	return w.roWALAdaptorImpl.Read(ctx, opts)
}

// retryAppendWhenRecoverableError retries the append operation when recoverable error occurs.
func (w *walAdaptorImpl) retryAppendWhenRecoverableError(ctx context.Context, msg message.MutableMessage) (message.MessageID, error) {
	backoff := backoff.NewExponentialBackOff()
	backoff.InitialInterval = 10 * time.Millisecond
	backoff.MaxInterval = 5 * time.Second
	backoff.MaxElapsedTime = 0
	backoff.Reset()

	// An append operation should be retried until it succeeds or some unrecoverable error occurs.
	for i := 0; ; i++ {
		appendCtx, span := message.StartSpanForMessage(ctx, msg, message.SpanNameWALAppendImpl)
		message.OverwriteTraceContext(appendCtx, msg)
		msgID, err := w.rwWALImpls.Append(appendCtx, msg)
		if err != nil {
			span.RecordError(err)
			span.SetStatus(codes.Error, err.Error())
		}
		span.End()
		if err == nil {
			if msg.MessageType() == message.MessageTypeAlterWAL {
				// if the append operation is a alter WAL message, we should log the message
				w.Logger().Info(ctx, "append alter WAL message to WAL finish", mlog.String("channel", msg.VChannel()), mlog.Uint64("timetick", msg.TimeTick()))
			}
			return msgID, nil
		}
		if errors.IsAny(err, context.Canceled, context.DeadlineExceeded, walimpls.ErrFenced) {
			return nil, err
		}
		w.writeMetrics.ObserveRetry()
		nextInterval := backoff.NextBackOff()
		w.Logger().Warn(ctx, "append message into wal impls failed, retrying...", mlog.FieldMessage(msg), mlog.Int("retry", i), mlog.Duration("nextInterval", nextInterval), mlog.Err(err))

		select {
		case <-ctx.Done():
			return nil, context.Cause(ctx)
		case <-w.availableCtx.Done():
			return nil, status.NewOnShutdownError("wal is on shutdown")
		case <-time.After(nextInterval):
		}
	}
}

// AppendAsync writes a record to the log asynchronously.
func (w *walAdaptorImpl) AppendAsync(ctx context.Context, msg message.MutableMessage, cb func(*wal.AppendResult, error)) {
	if !w.lifetime.Add(typeutil.LifetimeStateWorking) {
		cb(nil, status.NewOnShutdownError("wal is on shutdown"))
		return
	}

	// Submit async append to a background execution pool.
	_ = w.appendExecutionPool.Submit(func() (struct{}, error) {
		defer w.lifetime.Done()

		msgID, err := w.Append(ctx, msg)
		cb(msgID, err)
		return struct{}{}, nil
	})
}

// Close overrides Scanner Close function.
func (w *walAdaptorImpl) Close() {
	w.Logger().Info(context.TODO(), "wal begin to close, start graceful close...")
	// graceful close the interceptors before wal closing.
	w.interceptorBuildResult.GracefulCloseFunc()
	w.Logger().Info(context.TODO(), "wal graceful close done, wait for operation to be finished...")

	// begin to close the wal.
	w.lifetime.SetState(typeutil.LifetimeStateStopped)
	w.forceCancelAfterGracefulTimeout()
	w.lifetime.Wait()

	// close the flusher.
	w.Logger().Info(context.TODO(), "wal begin to close flusher...")
	if w.flusher != nil {
		// only in test, the flusher is nil.
		w.flusher.Close()
	}

	w.Logger().Info(context.TODO(), "wal begin to close scanners...")

	// close all wal instances.
	w.scanners.Range(func(id int64, s wal.Scanner) bool {
		s.Close()
		mlog.Info(context.TODO(), "close scanner by wal adaptor", mlog.Int64("id", id), mlog.Any("channel", w.Channel()))
		return true
	})

	w.Logger().Info(context.TODO(), "scanner close done, close inner wal...")
	w.rwWALImpls.Close()

	w.Logger().Info(context.TODO(), "wal close done, close interceptors...")
	w.interceptorBuildResult.Close()

	w.Logger().Info(context.TODO(), "close the write ahead buffer...")
	w.param.WriteAheadBuffer.Close()

	w.Logger().Info(context.TODO(), "close the segment assignment manager...")
	w.param.ShardManager.Close()

	w.Logger().Info(context.TODO(), "call wal cleanup function...")
	w.cleanup()
	w.Logger().Info(context.TODO(), "wal closed")

	// close all metrics.
	w.scanMetrics.Close()
	w.writeMetrics.Close()

	// close the rate limit component.
	w.WALRateLimitComponent.Close()

	if w.appendExecutionPool != nil {
		w.appendExecutionPool.Release()
	}
}

type interceptorBuildResult struct {
	Interceptor       interceptors.InterceptorWithReady
	GracefulCloseFunc gracefulCloseFunc
}

func (r interceptorBuildResult) Close() {
	r.Interceptor.Close()
}

// newWALWithInterceptors creates a new wal with interceptors.
func buildInterceptor(builders []interceptors.InterceptorBuilder, param *interceptors.InterceptorBuildParam) interceptorBuildResult {
	// Build all interceptors.
	builtIterceptors := make([]interceptors.Interceptor, 0, len(builders))
	for _, b := range builders {
		builtIterceptors = append(builtIterceptors, b.Build(param))
	}
	return interceptorBuildResult{
		Interceptor: interceptors.NewChainedInterceptor(builtIterceptors...),
		GracefulCloseFunc: func() {
			for _, i := range builtIterceptors {
				if c, ok := i.(interceptors.InterceptorWithGracefulClose); ok {
					c.GracefulClose()
				}
			}
		},
	}
}