milvus/internal/cdc/replication/replicatestream/replicate_stream_client_impl.go

// 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 replicatestream

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/cenkalti/backoff/v4"
	"github.com/cockroachdb/errors"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
	"github.com/milvus-io/milvus/internal/cdc/cluster"
	"github.com/milvus-io/milvus/internal/cdc/meta"
	"github.com/milvus-io/milvus/internal/cdc/resource"
	"github.com/milvus-io/milvus/internal/cdc/util"
	"github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil"
	"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/util/paramtable"
)

var ErrReplicationRemoved = errors.New("replication removed")

// replicateStreamClient is the implementation of ReplicateStreamClient.
type replicateStreamClient struct {
	clusterID       string
	targetClient    cluster.MilvusClient
	client          milvuspb.MilvusService_CreateReplicateStreamClient
	channel         *meta.ReplicateChannel
	pendingMessages MsgQueue
	metrics         ReplicateMetrics

	ctx        context.Context
	cancel     context.CancelFunc
	finishedCh chan struct{}
}

// NewReplicateStreamClient creates a new ReplicateStreamClient.
func NewReplicateStreamClient(ctx context.Context, c cluster.MilvusClient, channel *meta.ReplicateChannel) ReplicateStreamClient {
	ctx1, cancel := context.WithCancel(ctx)
	ctx1 = contextutil.WithClusterID(ctx1, channel.Value.GetTargetCluster().GetClusterId())

	options := MsgQueueOptions{
		Capacity: paramtable.Get().StreamingCfg.ReplicationPendingMessagesQueueLength.GetAsInt(),
		MaxSize:  paramtable.Get().StreamingCfg.ReplicationPendingMessagesQueueMaxSize.GetAsInt(),
	}
	pendingMessages := NewMsgQueue(options)
	rs := &replicateStreamClient{
		clusterID:       paramtable.Get().CommonCfg.ClusterPrefix.GetValue(),
		targetClient:    c,
		channel:         channel,
		pendingMessages: pendingMessages,
		metrics:         NewReplicateMetrics(channel.Value),
		ctx:             ctx1,
		cancel:          cancel,
		finishedCh:      make(chan struct{}),
	}

	rs.metrics.OnInitiate()
	go rs.startInternal()
	return rs
}

func (r *replicateStreamClient) startInternal() {
	defer func() {
		mlog.Info(r.ctx, "replicate stream client closed",
			mlog.String("key", r.channel.Key),
			mlog.Int64("revision", r.channel.ModRevision))
		r.metrics.OnClose()
		close(r.finishedCh)
	}()

	backoff := backoff.NewExponentialBackOff()
	backoff.InitialInterval = 100 * time.Millisecond
	backoff.MaxInterval = 10 * time.Second
	backoff.MaxElapsedTime = 0

	for {
		restart := r.startReplicating(backoff)
		if !restart {
			return
		}
		time.Sleep(backoff.NextBackOff())
	}
}

func (r *replicateStreamClient) startReplicating(backoff backoff.BackOff) (needRestart bool) {
	logger := mlog.With(mlog.String("key", r.channel.Key), mlog.Int64("revision", r.channel.ModRevision))
	if r.ctx.Err() != nil {
		logger.Info(r.ctx, "close replicate stream client due to ctx done")
		return false
	}

	// Create a local context for this connection that can be canceled
	// when we need to stop the send/recv loops
	connCtx, connCancel := context.WithCancel(r.ctx)
	defer connCancel()

	client, err := r.targetClient.CreateReplicateStream(connCtx)
	if err != nil {
		logger.Warn(r.ctx, "create milvus replicate stream failed, retry...", mlog.Err(err))
		return true
	}
	defer client.CloseSend()

	logger.Info(r.ctx, "replicate stream client service started")
	r.metrics.OnConnect()
	backoff.Reset()

	// reset client and pending messages
	r.client = client
	r.pendingMessages.SeekToHead()

	sendCh := r.startSendLoop(connCtx)
	recvCh := r.startRecvLoop(connCtx)

	var chErr error
	select {
	case <-r.ctx.Done():
	case chErr = <-sendCh:
	case chErr = <-recvCh:
	}

	connCancel() // Cancel the connection context
	<-sendCh
	<-recvCh // wait for send/recv loops to exit

	if r.ctx.Err() != nil {
		logger.Info(r.ctx, "close replicate stream client due to ctx done")
		return false
	} else if errors.Is(chErr, ErrReplicationRemoved) {
		logger.Info(r.ctx, "close replicate stream client due to replication removed")
		return false
	} else if isStreamIdleTimeout(chErr) {
		// Stream idle timeout is expected when no data is being replicated on the source channel.
		// See isStreamIdleTimeout for details.
		logger.Info(r.ctx, "replicate stream closed due to stream idle timeout, will reconnect", mlog.Err(chErr))
		r.metrics.OnDisconnect()
		return true
	} else {
		logger.Warn(r.ctx, "restart replicate stream client due to unexpected error", mlog.Err(chErr))
		r.metrics.OnDisconnect()
		return true
	}
}

// Replicate replicates the message to the target cluster.
func (r *replicateStreamClient) Replicate(msg message.ImmutableMessage) error {
	select {
	case <-r.ctx.Done():
		return nil
	default:
		if msg.MessageType().IsSelfControlled() {
			// If no messages are being replicated, update the last replicated time tick.
			if r.pendingMessages.Len() == 0 {
				r.metrics.UpdateLastReplicatedTimeTick(msg.TimeTick())
			}
			return ErrReplicateIgnored
		}
		r.metrics.StartReplicate(msg)
		r.pendingMessages.Enqueue(r.ctx, msg)
		return nil
	}
}

func (r *replicateStreamClient) startSendLoop(ctx context.Context) <-chan error {
	ch := make(chan error)
	go func() {
		err := r.sendLoop(ctx)
		ch <- err
		close(ch)
	}()
	return ch
}

func (r *replicateStreamClient) startRecvLoop(ctx context.Context) <-chan error {
	ch := make(chan error)
	go func() {
		err := r.recvLoop(ctx)
		ch <- err
		close(ch)
	}()
	return ch
}

func (r *replicateStreamClient) sendLoop(ctx context.Context) (err error) {
	logger := mlog.With(mlog.String("key", r.channel.Key), mlog.Int64("revision", r.channel.ModRevision))
	defer func() {
		if err != nil {
			logger.Warn(ctx, "send loop closed by unexpected error", mlog.Err(err))
		} else {
			logger.Info(ctx, "send loop closed")
		}
	}()
	for {
		select {
		case <-ctx.Done():
			return nil
		default:
			msg, err := r.pendingMessages.ReadNext(ctx)
			if err != nil {
				// context canceled, return nil
				return nil
			}
			if msg.MessageType() == message.MessageTypeTxn {
				txnMsg := message.AsImmutableTxnMessage(msg)
				err = r.sendTxnMessage(txnMsg)
				if err != nil {
					return err
				}
			} else {
				err = r.sendMessage(msg)
				if err != nil {
					return err
				}
			}
		}
	}
}

func (r *replicateStreamClient) sendMessage(msg message.ImmutableMessage) (err error) {
	immutableMessage := msg.IntoImmutableMessageProto()

	defer func() {
		logger := mlog.With(mlog.String("key", r.channel.Key), mlog.Int64("revision", r.channel.ModRevision))
		if err != nil {
			logger.Warn(r.ctx, "send message failed", mlog.Err(err), mlog.FieldMessage(msg))
		} else {
			r.metrics.OnSent(msg)
			logger.Debug(r.ctx, "send message success", mlog.FieldMessage(msg))
		}
	}()

	req := &milvuspb.ReplicateRequest{
		Request: &milvuspb.ReplicateRequest_ReplicateMessage{
			ReplicateMessage: &milvuspb.ReplicateMessage{
				SourceClusterId: r.clusterID,
				Message:         immutableMessage,
			},
		},
	}
	return r.client.Send(req)
}

func (r *replicateStreamClient) sendTxnMessage(txnMsg message.ImmutableTxnMessage) (err error) {
	// send txn begin message
	if err = r.sendMessage(txnMsg.Begin()); err != nil {
		return err
	}

	// send txn body messages
	if err = txnMsg.RangeOver(func(msg message.ImmutableMessage) error {
		return r.sendMessage(msg)
	}); err != nil {
		return err
	}

	// send txn commit message
	err = r.sendMessage(txnMsg.Commit())
	return
}

func (r *replicateStreamClient) recvLoop(ctx context.Context) (err error) {
	logger := mlog.With(mlog.String("key", r.channel.Key), mlog.Int64("revision", r.channel.ModRevision))
	defer func() {
		if err != nil && !errors.Is(err, ErrReplicationRemoved) {
			if isStreamIdleTimeout(err) {
				logger.Info(ctx, "replicate stream closed due to stream idle timeout, will reconnect", mlog.Err(err))
			} else {
				logger.Warn(ctx, "recv loop closed by unexpected error", mlog.Err(err))
			}
		} else {
			logger.Info(ctx, "recv loop closed", mlog.Err(err))
		}
	}()
	for {
		select {
		case <-ctx.Done():
			return nil
		default:
			resp, err := r.client.Recv()
			if err != nil {
				if isStreamIdleTimeout(err) {
					logger.Info(ctx, "replicate stream closed due to stream idle timeout, will reconnect", mlog.Err(err))
				} else {
					logger.Warn(ctx, "replicate stream recv failed", mlog.Err(err))
				}
				return err
			}
			lastConfirmedMessageInfo := resp.GetReplicateConfirmedMessageInfo()
			if lastConfirmedMessageInfo != nil {
				messages := r.pendingMessages.CleanupConfirmedMessages(lastConfirmedMessageInfo.GetConfirmedTimeTick())
				for _, msg := range messages {
					r.metrics.OnConfirmed(msg)
					if msg.MessageType() == message.MessageTypeAlterReplicateConfig {
						replicationRemoved := r.handleAlterReplicateConfigMessage(msg)
						if replicationRemoved {
							// Replication removed, return and stop replicate.
							return ErrReplicationRemoved
						}
					}
				}
			}
		}
	}
}

func (r *replicateStreamClient) handleAlterReplicateConfigMessage(msg message.ImmutableMessage) (replicationRemoved bool) {
	logger := mlog.With(mlog.String("key", r.channel.Key), mlog.Int64("revision", r.channel.ModRevision))
	logger.Info(r.ctx, "handle AlterReplicateConfigMessage", mlog.FieldMessage(msg))

	// Check ignore field - if true, skip processing
	// This is used for incomplete switchover messages that should be ignored after force promote
	alterMsg := message.MustAsImmutableAlterReplicateConfigMessageV2(msg)
	if alterMsg.Header().Ignore {
		logger.Info(r.ctx, "AlterReplicateConfig message has ignore flag set, skipping processing")
		return false
	}

	replicationRemoved = util.IsReplicationRemovedByAlterReplicateConfigMessage(msg, r.channel.Value)
	if replicationRemoved {
		// Cannot find the target channel, it means that the `current->target` topology edge is removed,
		// so we need to remove the replicate pchannel and stop replicate.
		etcdCli := resource.Resource().ETCD()
		ok, err := meta.RemoveReplicatePChannelWithRevision(r.ctx, etcdCli, r.channel.Key, r.channel.ModRevision)
		if err != nil {
			logger.Warn(r.ctx, "failed to remove replicate pchannel", mlog.Err(err))
			// When performing delete operation on etcd, the context may be canceled by the delete event
			// in cdc controller and then return `context.Canceled` error.
			// Since the delete event is generated after the delete operation is committed in etcd,
			// the delete is guaranteed to have succeeded on the server side.
			// So we can ignore the context canceled error here.
			if !errors.Is(err, context.Canceled) {
				panic(fmt.Sprintf("failed to remove replicate pchannel: %v", err))
			}
		}
		if ok {
			logger.Info(r.ctx, "handle AlterReplicateConfigMessage done, replicate pchannel removed")
		} else {
			logger.Info(r.ctx, "handle AlterReplicateConfigMessage done, revision not match, replicate pchannel not removed")
		}
		return true
	}
	logger.Info(r.ctx, "target channel found, skip handle AlterReplicateConfigMessage")
	return false
}

func (r *replicateStreamClient) BlockUntilFinish() {
	<-r.finishedCh
}

func (r *replicateStreamClient) Close() {
	r.cancel()
	<-r.finishedCh
}

// isStreamIdleTimeout checks if the error is a gRPC "stream timeout" error.
// This is typically caused by envoy sidecar's stream_idle_timeout (default 5m):
// when no application-level DATA frames flow on a gRPC bidirectional stream,
// envoy considers the stream idle and terminates it, even though gRPC transport-level
// keepalive pings are still active (PING frames don't reset stream_idle_timeout).
// This is expected when no data is being replicated on the source channel.
func isStreamIdleTimeout(err error) bool {
	if err == nil {
		return false
	}
	s, ok := status.FromError(err)
	return ok && s.Code() == codes.Unknown && strings.Contains(s.Message(), "stream timeout")
}