internal/topic/topicreaderinternal/stream_reader_impl.go - master - githubmirror/ydb-go-sdk

internal/topic/topicreaderinternal/stream_reader_impl.go
1 048 строк29 KB
Aleksey Myasnikov
Moved the trace-generated code in `trace` package from the public API to the internal packages (#2172)
24 май 2026, 16:46
24 май 2026, 16:4686fc0c2
package topicreaderinternal

import (
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"math"
	"math/big"
	"reflect"
	"runtime/pprof"
	"sync/atomic"
	"time"

	"github.com/ydb-platform/ydb-go-sdk/v3/credentials"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/background"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopicreader"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/gtrace"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicreadercommon"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
)

var (
	errCommitWithNilPartitionSession = xerrors.Wrap(errors.New("ydb: commit with nil partition session"))
	errUnexpectedEmptyConsumerName   = xerrors.Wrap(errors.New("ydb: create ydb reader with empty consumer name. Set one of: consumer name or option WithReaderWithoutConsumer")) //nolint:lll
	errCantCommitWithoutConsumer     = xerrors.Wrap(errors.New("ydb: reader can't commit messages without consumer"))
	errBufferSize                    = xerrors.Wrap(errors.New("ydb: buffer of topic reader must be greater than zero, see option topicoptions.WithReaderBufferSizeBytes")) //nolint:lll
	errTopicSelectorsEmpty           = xerrors.Wrap(errors.New("ydb: topic selector for topic reader is empty, see arguments on topic starts"))                             //nolint:lll
)

var clientSessionCounter atomic.Int64

type partitionSessionID = rawtopicreader.PartitionSessionID

type topicStreamReaderImpl struct {
	cfg    topicStreamReaderConfig
	ctx    context.Context //nolint:containedctx
	cancel context.CancelFunc

	topicClient         TopicClient
	freeBytes           chan int
	restBufferSizeBytes atomic.Int64
	sessionController   topicreadercommon.PartitionSessionStorage
	backgroundWorkers   background.Worker

	rawMessagesFromBuffer chan rawtopicreader.ServerMessage

	batcher   *batcher
	committer *topicreadercommon.Committer

	stream           topicreadercommon.RawTopicReaderStream
	readConnectionID string
	readerID         int64

	m       xsync.RWMutex
	err     error
	started bool
	closed  bool

	batchTxStorage *batchTxStorage
}

type topicStreamReaderConfig struct {
	CommitterBatchTimeLag           time.Duration
	CommitterBatchCounterTrigger    int
	BaseContext                     context.Context //nolint:containedctx
	BufferSizeProtoBytes            int
	Cred                            credentials.Credentials
	CredUpdateInterval              time.Duration
	Consumer                        string
	ReadWithoutConsumer             bool
	ReadSelectors                   []*topicreadercommon.PublicReadSelector
	Trace                           *trace.Topic
	GetPartitionStartOffsetCallback PublicGetPartitionStartOffsetFunc
	OnStopPartitionSession          PublicOnStopPartitionSessionFunc
	CommitMode                      topicreadercommon.PublicCommitMode
	Decoders                        *topicreadercommon.MultiDecoder
	EnableSplitMergeSupport         bool
}

func newTopicStreamReaderConfig() topicStreamReaderConfig {
	return topicStreamReaderConfig{
		BaseContext:             context.Background(),
		BufferSizeProtoBytes:    topicreadercommon.DefaultBufferSize,
		Cred:                    credentials.NewAnonymousCredentials(),
		CredUpdateInterval:      time.Hour,
		CommitMode:              topicreadercommon.CommitModeAsync,
		CommitterBatchTimeLag:   time.Second,
		Decoders:                topicreadercommon.NewMultiDecoder(),
		Trace:                   &trace.Topic{},
		EnableSplitMergeSupport: true,
	}
}

func (cfg *topicStreamReaderConfig) Validate() []error {
	var validateErrors []error

	if cfg.Consumer != "" && cfg.ReadWithoutConsumer {
		validateErrors = append(validateErrors, errSetConsumerAndNoConsumer)
	}
	if cfg.Consumer == "" && !cfg.ReadWithoutConsumer {
		validateErrors = append(validateErrors, errUnexpectedEmptyConsumerName)
	}
	if cfg.ReadWithoutConsumer && cfg.CommitMode != topicreadercommon.CommitModeNone {
		validateErrors = append(validateErrors, errCantCommitWithoutConsumer)
	}
	if cfg.BufferSizeProtoBytes <= 0 {
		validateErrors = append(validateErrors, errBufferSize)
	}
	if len(cfg.ReadSelectors) == 0 {
		validateErrors = append(validateErrors, errTopicSelectorsEmpty)
	}

	return validateErrors
}

func newTopicStreamReader(
	client TopicClient,
	readerID int64,
	stream topicreadercommon.RawTopicReaderStream,
	cfg topicStreamReaderConfig, //nolint:gocritic
) (_ *topicStreamReaderImpl, err error) {
	defer func() {
		if err != nil {
			_ = stream.CloseSend()
		}
	}()

	reader := newTopicStreamReaderStopped(client, readerID, stream, cfg)
	if err = reader.initSession(); err != nil {
		return nil, err
	}
	if err = reader.startBackgroundWorkers(); err != nil {
		return nil, err
	}

	return reader, nil
}

func newTopicStreamReaderStopped(
	client TopicClient,
	readerID int64,
	stream topicreadercommon.RawTopicReaderStream,
	cfg topicStreamReaderConfig, //nolint:gocritic
) *topicStreamReaderImpl {
	labeledContext := pprof.WithLabels(cfg.BaseContext, pprof.Labels("base-context", "topic-stream-reader"))
	stopPump, cancel := xcontext.WithCancel(labeledContext)

	readerConnectionID, err := rand.Int(rand.Reader, big.NewInt(math.MaxInt64))
	if err != nil {
		readerConnectionID = big.NewInt(-1)
	}

	res := &topicStreamReaderImpl{
		cfg:                   cfg,
		ctx:                   stopPump,
		topicClient:           client,
		freeBytes:             make(chan int, 1),
		stream:                topicreadercommon.NewSyncedStream(stream),
		cancel:                cancel,
		batcher:               newBatcher(),
		readConnectionID:      "preinitID-" + readerConnectionID.String(),
		readerID:              readerID,
		rawMessagesFromBuffer: make(chan rawtopicreader.ServerMessage, 1),
	}

	res.backgroundWorkers = *background.NewWorker(stopPump, "topic-reader-stream-background")

	res.committer = topicreadercommon.NewCommitterStopped(cfg.Trace, labeledContext, cfg.CommitMode, res.send)
	res.committer.BufferTimeLagTrigger = cfg.CommitterBatchTimeLag
	res.committer.BufferCountTrigger = cfg.CommitterBatchCounterTrigger
	res.freeBytes <- cfg.BufferSizeProtoBytes

	res.batchTxStorage = newBatchTxStorage(cfg.Consumer)

	return res
}

func (r *topicStreamReaderImpl) ReadSessionID() string {
	return r.readConnectionID
}

func (r *topicStreamReaderImpl) WaitInit(_ context.Context) error {
	if !r.started {
		return errors.New("not started: can be started only after initialize from constructor")
	}

	return nil
}

func (r *topicStreamReaderImpl) TopicOnReaderStart(consumer string, err error) {
	logCtx := r.cfg.BaseContext
	gtrace.TopicOnReaderStart(r.cfg.Trace, &logCtx, r.readerID, consumer, err)
}

func (r *topicStreamReaderImpl) PopMessagesBatchTx(
	ctx context.Context,
	tx tx.Transaction,
	opts ReadMessageBatchOptions,
) (_ *topicreadercommon.PublicBatch, resErr error) {
	logCtx := r.cfg.BaseContext
	onDone := gtrace.TopicOnReaderStreamPopBatchTx(
		r.cfg.Trace,
		&logCtx,
		r.readerID,
		r.readConnectionID,
		tx.SessionID(),
		tx,
	)
	defer func() {
		onDone(resErr)
	}()

	batch, err := r.ReadMessageBatch(ctx, opts)
	if err != nil {
		return nil, err
	}

	if err = r.commitWithTransaction(ctx, tx, batch); err == nil {
		return batch, nil
	}

	return nil, err
}

// commitWithTransaction prepares a batch of messages to be committed within a YDB transaction.
//
// This method does not immediately commit the batch. Instead, it registers hooks that will be
// executed when the transaction is committed:
//  1. Materializes the transaction
//  2. Registers a pre-commit hook to update topic offsets
//  3. Registers a completion hook to update local offsets or close the reader on failure
//
// The actual batch commit happens when the transaction is committed and the registered hooks are executed.
// The offset update is performed on the same node where the transaction was initiated to avoid
// coordination errors. All errors from UpdateOffsetsInTransaction are treated as retryable.
//
// Parameters:
//   - ctx: context for the operation
//   - tx: the YDB transaction to associate the batch commit with
//   - batch: the batch of messages to be committed when the transaction commits
//
// Returns an error if transaction materialization fails. Offset update errors are handled
// within the transaction's pre-commit hook and will cause the transaction to fail.
func (r *topicStreamReaderImpl) commitWithTransaction(
	ctx context.Context,
	tx tx.Transaction,
	batch *topicreadercommon.PublicBatch,
) error {
	if err := tx.UnLazy(ctx); err != nil {
		return fmt.Errorf("ydb: failed to materialize transaction: %w", err)
	}

	txBatches, isNew := r.batchTxStorage.GetOrCreateTransactionBatches(tx)
	txBatches.AddBatch(batch)
	if !isNew {
		// tx hooks already configured - exiting
		return nil
	}

	tx.OnBeforeCommit(r.txBeforeCommitFn(tx))
	tx.OnCompleted(r.txOnCompletedFn(ctx, tx))

	return nil
}

func (r *topicStreamReaderImpl) txBeforeCommitFn(tx tx.Transaction) tx.OnTransactionBeforeCommit {
	return func(ctx context.Context) (err error) {
		logCtx := r.cfg.BaseContext
		onDone := gtrace.TopicOnReaderUpdateOffsetsInTransaction(
			r.cfg.Trace,
			&logCtx,
			r.readerID,
			r.readConnectionID,
			tx.SessionID(),
			tx,
		)
		defer func() {
			onDone(err)
		}()

		// UpdateOffsetsInTransaction operation must be executed on the same Node where the transaction was initiated.
		// Otherwise, errors such as `Database coordinators are unavailable` may occur.
		ctx = endpoint.WithNodeID(ctx, tx.NodeID())

		req, err := r.batchTxStorage.GetUpdateOffsetsInTransactionRequest(tx)
		if err != nil {
			if errors.Is(err, errNoBatches) {
				return nil
			}

			return xerrors.WithStackTrace(fmt.Errorf("building update offsets request: %w", err))
		}

		err = r.topicClient.UpdateOffsetsInTransaction(ctx, req)
		if err != nil {
			return xerrors.WithStackTrace(fmt.Errorf("updating offsets in transaction: %w", err))
		}

		return nil
	}
}

func (r *topicStreamReaderImpl) txOnCompletedFn(ctx context.Context, tx tx.Transaction) tx.OnTransactionCompletedFunc {
	return func(err error) {
		logCtx := r.cfg.BaseContext
		onDone := gtrace.TopicOnReaderTransactionCompleted(
			r.cfg.Trace,
			&logCtx,
			r.readerID,
			r.readConnectionID,
			tx.SessionID(),
			tx,
			err,
		)
		defer onDone()

		defer r.batchTxStorage.Clear(tx)

		if err != nil {
			// mark error as retryable - for proper reconnector working
			err = xerrors.Retryable(err)

			_ = r.CloseWithError(xcontext.ValueOnly(ctx), fmt.Errorf("transaction failed: %w", err))

			return
		}

		for _, batch := range r.batchTxStorage.GetBatches(tx) {
			commitRange := topicreadercommon.GetCommitRange(batch)
			topicreadercommon.BatchGetPartitionSession(batch).SetCommittedOffsetForward(commitRange.CommitOffsetEnd)
		}
	}
}

func (r *topicStreamReaderImpl) ReadMessageBatch(
	ctx context.Context,
	opts ReadMessageBatchOptions,
) (batch *topicreadercommon.PublicBatch, err error) {
	mergeCtx := xcontext.MergeContexts(ctx, r.cfg.BaseContext)
	onDone := gtrace.TopicOnReaderReadMessages(
		r.cfg.Trace,
		&mergeCtx,
		opts.MinCount,
		opts.MaxCount,
		r.getRestBufferBytes(),
	)
	defer func() {
		if batch == nil {
			onDone(0, "", -1, -1, -1, -1, r.getRestBufferBytes(), err)
		} else {
			commitRange := topicreadercommon.GetCommitRange(batch)
			onDone(
				len(batch.Messages),
				batch.Topic(),
				batch.PartitionID(),
				topicreadercommon.BatchGetPartitionSession(batch).StreamPartitionSessionID.ToInt64(),
				commitRange.CommitOffsetStart.ToInt64(),
				commitRange.CommitOffsetEnd.ToInt64(),
				r.getRestBufferBytes(),
				err,
			)
		}
	}()

	if err = ctx.Err(); err != nil {
		return nil, err
	}

	defer func() {
		if err == nil {
			r.freeBufferFromMessages(batch)
		}
	}()

	return r.consumeMessagesUntilBatch(ctx, opts)
}

func (r *topicStreamReaderImpl) consumeMessagesUntilBatch(
	ctx context.Context,
	opts ReadMessageBatchOptions,
) (*topicreadercommon.PublicBatch, error) {
	for {
		item, err := r.batcher.Pop(ctx, opts.batcherGetOptions)
		if err != nil {
			return nil, err
		}

		switch {
		case item.IsBatch():
			return item.Batch, nil
		case item.IsRawMessage():
			r.sendRawMessageToChannelUnblocked(item.RawMessage)
		default:
			return nil, xerrors.WithStackTrace(fmt.Errorf("ydb: unexpected item type from batcher: %#v", item))
		}
	}
}

func (r *topicStreamReaderImpl) sendRawMessageToChannelUnblocked(msg rawtopicreader.ServerMessage) {
	select {
	case r.rawMessagesFromBuffer <- msg:
		return
	default:
		// send in goroutine, without block caller
		r.backgroundWorkers.Start("sendMessageToRawChannel", func(ctx context.Context) {
			select {
			case r.rawMessagesFromBuffer <- msg:
			case <-ctx.Done():
			}
		})
	}
}

func (r *topicStreamReaderImpl) consumeRawMessageFromBuffer(ctx context.Context) {
	doneChan := ctx.Done()

	for {
		var msg rawtopicreader.ServerMessage
		select {
		case <-doneChan:
			return
		case msg = <-r.rawMessagesFromBuffer:
			// pass
		}

		switch m := msg.(type) {
		case *rawtopicreader.StartPartitionSessionRequest:
			if err := r.onStartPartitionSessionRequestFromBuffer(m); err != nil {
				_ = r.CloseWithError(ctx, err)

				return
			}
		case *rawtopicreader.StopPartitionSessionRequest:
			if err := r.onStopPartitionSessionRequestFromBuffer(m); err != nil {
				_ = r.CloseWithError(ctx, xerrors.WithStackTrace(
					fmt.Errorf("ydb: unexpected error on stop partition handler: %w", err),
				))

				return
			}
		case *rawtopicreader.PartitionSessionStatusResponse:
			r.onPartitionSessionStatusResponseFromBuffer(ctx, m)
		default:
			_ = r.CloseWithError(ctx, xerrors.WithStackTrace(
				fmt.Errorf("ydb: unexpected server message from buffer: %v", reflect.TypeOf(msg))),
			)
		}
	}
}

func (r *topicStreamReaderImpl) onStopPartitionSessionRequestFromBuffer(
	msg *rawtopicreader.StopPartitionSessionRequest,
) (err error) {
	session, err := r.sessionController.Get(msg.PartitionSessionID)
	if err != nil {
		return err
	}
	onDone := gtrace.TopicOnReaderPartitionReadStopResponse(
		r.cfg.Trace,
		r.readConnectionID,
		session.Context(),
		session.Topic,
		session.PartitionID,
		session.StreamPartitionSessionID.ToInt64(),
		msg.CommittedOffset.ToInt64(),
		msg.Graceful,
	)
	defer func() {
		onDone(err)
	}()

	if r.cfg.OnStopPartitionSession != nil {
		_ = r.cfg.OnStopPartitionSession(PublicStopPartitionSessionRequest{
			Topic:              session.Topic,
			PartitionID:        session.PartitionID,
			PartitionSessionID: session.StreamPartitionSessionID.ToInt64(),
			CommittedOffset:    msg.CommittedOffset.ToInt64(),
			Graceful:           msg.Graceful,
		})
	}

	if msg.Graceful {
		if err = r.committer.Flush(); err != nil {
			return xerrors.WithStackTrace(fmt.Errorf(
				"ydb: flush buffered commits on graceful stop partition session: %w",
				err,
			))
		}
		session.Close()
		resp := &rawtopicreader.StopPartitionSessionResponse{
			PartitionSessionID: session.StreamPartitionSessionID,
		}
		if err = r.send(resp); err != nil {
			return err
		}
	}

	if _, err = r.sessionController.Remove(session.StreamPartitionSessionID); err != nil {
		if msg.Graceful {
			return err
		} else { //nolint:revive,staticcheck
			// double message with graceful=false is ok.
			// It may be received after message with graceful=true and session was removed while process that.

			// pass
		}
	}

	return nil
}

func (r *topicStreamReaderImpl) onPartitionSessionStatusResponseFromBuffer(
	ctx context.Context,
	m *rawtopicreader.PartitionSessionStatusResponse,
) {
	panic("not implemented")
}

func (r *topicStreamReaderImpl) onUpdateTokenResponse(m *rawtopicreader.UpdateTokenResponse) {
}

func (r *topicStreamReaderImpl) Commit(ctx context.Context, commitRange topicreadercommon.CommitRange) (err error) {
	defer func() {
		if errors.Is(
			err,
			topicreadercommon.ErrPublicCommitSessionToExpiredSession,
		) && r.cfg.CommitMode == topicreadercommon.CommitModeAsync {
			err = nil
		}
	}()

	if commitRange.PartitionSession == nil {
		return xerrors.WithStackTrace(errCommitWithNilPartitionSession)
	}

	session := commitRange.PartitionSession
	mergeCtx := xcontext.MergeContexts(ctx, r.cfg.BaseContext)
	onDone := gtrace.TopicOnReaderCommit(
		r.cfg.Trace,
		&mergeCtx,
		session.Topic,
		session.PartitionID,
		session.StreamPartitionSessionID.ToInt64(),
		commitRange.CommitOffsetStart.ToInt64(),
		commitRange.CommitOffsetEnd.ToInt64(),
	)

	defer func() {
		onDone(err)
	}()

	if err = r.checkCommitRange(commitRange); err != nil {
		return err
	}

	return r.committer.Commit(ctx, commitRange)
}

func (r *topicStreamReaderImpl) checkCommitRange(commitRange topicreadercommon.CommitRange) error {
	if r.cfg.CommitMode == topicreadercommon.CommitModeNone {
		return topicreadercommon.ErrCommitDisabled
	}
	session := commitRange.PartitionSession

	if session == nil {
		return xerrors.WithStackTrace(errCommitWithNilPartitionSession)
	}

	if session.Context().Err() != nil {
		return xerrors.WithStackTrace(topicreadercommon.ErrPublicCommitSessionToExpiredSession)
	}

	ownSession, err := r.sessionController.Get(session.StreamPartitionSessionID)
	if err != nil || session != ownSession {
		return xerrors.WithStackTrace(topicreadercommon.ErrPublicCommitSessionToExpiredSession)
	}
	if session.CommittedOffset() != commitRange.CommitOffsetStart && r.cfg.CommitMode == topicreadercommon.CommitModeSync {
		return topicreadercommon.ErrWrongCommitOrderInSyncMode
	}

	return nil
}

func (r *topicStreamReaderImpl) send(msg rawtopicreader.ClientMessage) error {
	err := r.stream.Send(msg)
	if err != nil {
		logCtx := r.cfg.BaseContext
		gtrace.TopicOnReaderError(r.cfg.Trace, &logCtx, r.readConnectionID, err)

		_ = r.CloseWithError(r.ctx, err)
	}

	return err
}

func (r *topicStreamReaderImpl) startBackgroundWorkers() error {
	if err := r.setStarted(); err != nil {
		return err
	}

	r.committer.Start()

	r.backgroundWorkers.Start("readMessagesLoop", r.readMessagesLoop)
	r.backgroundWorkers.Start("dataRequestLoop", r.dataRequestLoop)
	r.backgroundWorkers.Start("updateTokenLoop", r.updateTokenLoop)

	r.backgroundWorkers.Start("consumeRawMessageFromBuffer", r.consumeRawMessageFromBuffer)

	return nil
}

func (r *topicStreamReaderImpl) setStarted() error {
	r.m.Lock()
	defer r.m.Unlock()

	if r.started {
		return xerrors.WithStackTrace(errors.New("already started"))
	}

	r.started = true

	return nil
}

func (r *topicStreamReaderImpl) initSession() (err error) {
	initMessage := topicreadercommon.CreateInitMessage(r.cfg.Consumer, r.cfg.EnableSplitMergeSupport, r.cfg.ReadSelectors)

	logCtx := r.cfg.BaseContext
	onDone := gtrace.TopicOnReaderInit(r.cfg.Trace, &logCtx, r.readConnectionID, initMessage)

	defer func() {
		onDone(r.readConnectionID, err)
	}()

	if err = r.send(initMessage); err != nil {
		return err
	}

	resp, err := r.stream.Recv()
	if err != nil {
		return err
	}

	if status := resp.StatusData(); !status.Status.IsSuccess() {
		// Need wrap status to common ydb operational error
		// https://github.com/ydb-platform/ydb-go-sdk/issues/1361
		return xerrors.WithStackTrace(fmt.Errorf("bad status on initial error: %v (%v)", status.Status, status.Issues))
	}

	initResp, ok := resp.(*rawtopicreader.InitResponse)
	if !ok {
		return xerrors.WithStackTrace(fmt.Errorf("bad message type on session init: %v (%v)", resp, reflect.TypeOf(resp)))
	}

	r.readConnectionID = initResp.SessionID

	return nil
}

func (r *topicStreamReaderImpl) addRestBufferBytes(delta int) int {
	val := r.restBufferSizeBytes.Add(int64(delta))
	if val <= 0 {
		r.batcher.IgnoreMinRestrictionsOnNextPop()
	}

	return int(val)
}

func (r *topicStreamReaderImpl) getRestBufferBytes() int {
	return int(r.restBufferSizeBytes.Load())
}

//nolint:funlen
func (r *topicStreamReaderImpl) readMessagesLoop(ctx context.Context) {
	ctx, cancel := xcontext.WithCancel(ctx)
	defer cancel()

	for {
		serverMessage, err := r.stream.Recv()
		if err != nil {
			logCtx := r.cfg.BaseContext
			gtrace.TopicOnReaderError(r.cfg.Trace, &logCtx, r.readConnectionID, err)
			if errors.Is(err, rawtopicreader.ErrUnexpectedMessageType) {
				gtrace.TopicOnReaderUnknownGrpcMessage(r.cfg.Trace, &logCtx, r.readConnectionID, err)
				// new messages can be added to protocol, it must be backward compatible to old programs
				// and skip message is safe
				continue
			}
			_ = r.CloseWithError(ctx, err)

			return
		}

		status := serverMessage.StatusData()
		if !status.Status.IsSuccess() {
			_ = r.CloseWithError(ctx,
				xerrors.WithStackTrace(
					fmt.Errorf("ydb: bad status from pq grpc stream: %v, %v", status.Status, status.Issues.String()),
				),
			)
		}

		switch m := serverMessage.(type) {
		case *rawtopicreader.ReadResponse:
			if err = r.onReadResponse(m); err != nil {
				_ = r.CloseWithError(ctx, err)
			}
		case *rawtopicreader.StartPartitionSessionRequest:
			if err = r.onStartPartitionSessionRequest(m); err != nil {
				_ = r.CloseWithError(ctx, err)

				return
			}
		case *rawtopicreader.StopPartitionSessionRequest:
			if err = r.onStopPartitionSessionRequest(m); err != nil {
				_ = r.CloseWithError(ctx, err)

				return
			}
		case *rawtopicreader.EndPartitionSession:
			if err = r.onEndPartitionSession(m); err != nil {
				_ = r.CloseWithError(ctx, err)
			}
		case *rawtopicreader.CommitOffsetResponse:
			if err = r.onCommitResponse(m); err != nil {
				_ = r.CloseWithError(ctx, err)

				return
			}

		case *rawtopicreader.UpdateTokenResponse:
			r.onUpdateTokenResponse(m)
		default:
			{
				logCtx := r.cfg.BaseContext
				gtrace.TopicOnReaderUnknownGrpcMessage(
					r.cfg.Trace,
					&logCtx,
					r.readConnectionID,
					xerrors.WithStackTrace(xerrors.Wrap(fmt.Errorf(
						"ydb: unexpected message type in stream reader: %v",
						reflect.TypeOf(serverMessage),
					))),
				)
			}
		}
	}
}

func (r *topicStreamReaderImpl) dataRequestLoop(ctx context.Context) {
	if r.ctx.Err() != nil {
		return
	}

	doneChan := ctx.Done()

	for {
		select {
		case <-doneChan:
			_ = r.CloseWithError(ctx, r.ctx.Err())

			return

		case free := <-r.freeBytes:
			sum := free

			// consume all messages from order and compress it to one data request
		forConsumeRequests:
			for {
				select {
				case free = <-r.freeBytes:
					sum += free
				default:
					break forConsumeRequests
				}
			}

			resCapacity := r.addRestBufferBytes(sum)
			logCtx := r.cfg.BaseContext
			gtrace.TopicOnReaderSentDataRequest(r.cfg.Trace, &logCtx, r.readConnectionID, sum, resCapacity)
			if err := r.sendDataRequest(sum); err != nil {
				return
			}
		}
	}
}

func (r *topicStreamReaderImpl) sendDataRequest(size int) error {
	return r.send(&rawtopicreader.ReadRequest{BytesSize: size})
}

func (r *topicStreamReaderImpl) freeBufferFromMessages(batch *topicreadercommon.PublicBatch) {
	size := 0
	for messageIndex := range batch.Messages {
		size += topicreadercommon.MessageGetBufferBytesAccount(batch.Messages[messageIndex])
	}
	select {
	case r.freeBytes <- size:
	case <-r.ctx.Done():
	}
}

func (r *topicStreamReaderImpl) updateTokenLoop(ctx context.Context) {
	ticker := time.NewTicker(r.cfg.CredUpdateInterval)
	defer ticker.Stop()

	readerCancel := ctx.Done()
	for {
		select {
		case <-readerCancel:
			return
		case <-ticker.C:
			r.updateToken(r.ctx)
		}
	}
}

func (r *topicStreamReaderImpl) onReadResponse(msg *rawtopicreader.ReadResponse) (err error) {
	resCapacity := r.addRestBufferBytes(-msg.BytesSize)

	logCtx := r.cfg.BaseContext
	onDone := gtrace.TopicOnReaderReceiveDataResponse(r.cfg.Trace, &logCtx, r.readConnectionID, resCapacity, msg)
	defer func() {
		onDone(err)
	}()

	batches, err2 := topicreadercommon.ReadRawBatchesToPublicBatches(msg, &r.sessionController, r.cfg.Decoders)
	if err2 != nil {
		return err2
	}

	for i := range batches {
		if err := r.batcher.PushBatches(batches[i]); err != nil {
			return err
		}
	}

	return nil
}

func (r *topicStreamReaderImpl) CloseWithError(ctx context.Context, reason error) (closeErr error) {
	logCtx := r.cfg.BaseContext
	onDone := gtrace.TopicOnReaderClose(r.cfg.Trace, &logCtx, r.readConnectionID, reason)
	defer func() {
		onDone(closeErr)
	}()

	isFirstClose := false
	r.m.WithLock(func() {
		if r.closed {
			return
		}
		isFirstClose = true
		r.closed = true

		r.err = reason
		r.cancel()
	})
	if !isFirstClose {
		return nil
	}

	closeErr = r.committer.Close(ctx, reason)

	batcherErr := r.batcher.Close(reason)
	if closeErr == nil {
		closeErr = batcherErr
	}

	// close stream strong after committer close - for flush commits buffer
	streamCloseErr := r.stream.CloseSend()
	if closeErr == nil {
		closeErr = streamCloseErr
	}

	// close background workers after r.stream.CloseSend
	bgCloseErr := r.backgroundWorkers.Close(ctx, reason)
	if closeErr == nil {
		closeErr = bgCloseErr
	}

	return closeErr
}

func (r *topicStreamReaderImpl) onCommitResponse(msg *rawtopicreader.CommitOffsetResponse) error {
	for i := range msg.PartitionsCommittedOffsets {
		commit := &msg.PartitionsCommittedOffsets[i]
		partition, err := r.sessionController.Get(commit.PartitionSessionID)
		if err != nil {
			return fmt.Errorf("ydb: can't found session on commit response: %w", err)
		}
		partition.SetCommittedOffsetForward(commit.CommittedOffset)

		logCtx := r.cfg.BaseContext
		gtrace.TopicOnReaderCommittedNotify(
			r.cfg.Trace,
			&logCtx,
			r.readConnectionID,
			partition.Topic,
			partition.PartitionID,
			partition.StreamPartitionSessionID.ToInt64(),
			commit.CommittedOffset.ToInt64(),
		)

		r.committer.OnCommitNotify(partition, commit.CommittedOffset)
	}

	return nil
}

func (r *topicStreamReaderImpl) updateToken(ctx context.Context) {
	logCtx := r.cfg.BaseContext
	onUpdateToken := gtrace.TopicOnReaderUpdateToken(
		r.cfg.Trace,
		&logCtx,
		r.readConnectionID,
	)
	token, err := r.cfg.Cred.Token(ctx)
	onSent := onUpdateToken(&ctx, len(token), err)
	if err != nil {
		return
	}

	err = r.send(&rawtopicreader.UpdateTokenRequest{UpdateTokenRequest: rawtopiccommon.UpdateTokenRequest{Token: token}})
	onSent(err)
}

func (r *topicStreamReaderImpl) onStartPartitionSessionRequest(m *rawtopicreader.StartPartitionSessionRequest) error {
	session := topicreadercommon.NewPartitionSession(
		r.ctx,
		m.PartitionSession.Path,
		m.PartitionSession.PartitionID,
		r.readerID,
		r.readConnectionID,
		m.PartitionSession.PartitionSessionID,
		clientSessionCounter.Add(1),
		m.CommittedOffset,
	)
	if err := r.sessionController.Add(session); err != nil {
		return err
	}

	return r.batcher.PushRawMessage(session, m)
}

func (r *topicStreamReaderImpl) onStartPartitionSessionRequestFromBuffer(
	m *rawtopicreader.StartPartitionSessionRequest,
) (err error) {
	session, err := r.sessionController.Get(m.PartitionSession.PartitionSessionID)
	if err != nil {
		return err
	}

	var (
		ctx    = session.Context()
		onDone = gtrace.TopicOnReaderPartitionReadStartResponse(
			r.cfg.Trace,
			r.readConnectionID,
			&ctx,
			session.Topic,
			session.PartitionID,
			session.StreamPartitionSessionID.ToInt64(),
		)
	)

	respMessage := &rawtopicreader.StartPartitionSessionResponse{
		PartitionSessionID: session.StreamPartitionSessionID,
	}

	var forceOffset *int64
	var commitOffset *int64

	defer func() {
		onDone(forceOffset, commitOffset, err)
	}()

	if r.cfg.GetPartitionStartOffsetCallback != nil {
		req := PublicGetPartitionStartOffsetRequest{
			Topic:       session.Topic,
			PartitionID: session.PartitionID,
		}
		resp, callbackErr := r.cfg.GetPartitionStartOffsetCallback(session.Context(), req)
		if callbackErr != nil {
			return callbackErr
		}
		if resp.startOffsetUsed {
			wantOffset := resp.startOffset.ToInt64()
			forceOffset = &wantOffset
		}
	}

	respMessage.ReadOffset.FromInt64Pointer(forceOffset)
	if r.cfg.CommitMode.CommitsEnabled() {
		commitOffset = forceOffset
		respMessage.CommitOffset.FromInt64Pointer(commitOffset)
	}
	if forceOffset != nil {
		session.SetInitialCommitOffset(rawtopiccommon.NewOffset(*forceOffset))
	}

	return r.send(respMessage)
}

func (r *topicStreamReaderImpl) onStopPartitionSessionRequest(m *rawtopicreader.StopPartitionSessionRequest) error {
	session, err := r.sessionController.Get(m.PartitionSessionID)
	if err != nil {
		return err
	}

	session.SetNoMoreMessages()
	if !m.Graceful {
		session.Close()
	}

	return r.batcher.PushRawMessage(session, m)
}

func (r *topicStreamReaderImpl) onEndPartitionSession(m *rawtopicreader.EndPartitionSession) error {
	// need err value in else block
	//nolint:revive
	if session, err := r.sessionController.Get(m.PartitionSessionID); err == nil {
		gtrace.TopicOnReaderEndPartitionSession(
			r.cfg.Trace,
			r.readConnectionID,
			session.Context(),
			session.Topic,
			session.PartitionID,
			m.PartitionSessionID.ToInt64(),
			m.AdjacentPartitionIDs,
			m.ChildPartitionIDs,
		)
		session.SetNoMoreMessages()
		r.batcher.FlushPartitionSession(session)

		return nil
	} else {
		return xerrors.Retryable(xerrors.Wrap(fmt.Errorf(
			"ydb: unknown partition for end partition session: %w", err,
		)))
	}
}