Merge pull request #1358 Topic callback reader initial experimental

Author: rekby

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Added experimental topic listener implementation
 * Fixed `internal/xstrings.Buffer()` leak without call `buffer.Free()`
 * Removed double quotas from goroutine labels background workers for prevent problem with pprof
 

examples/topic/topicreader/topicreader_handler.go

@@ -0,0 +1,97 @@
+package topicreaderexamples
+
+import (
+	"context"
+	"log"
+	"sync"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3"
+	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topiclistener"
+	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicoptions"
+)
+
+func StartReader(ctx context.Context, db *ydb.Driver) (*topiclistener.TopicListener, error) {
+	handler := &TopicEventsHandler{
+		locks: make(map[int64]int64),
+	}
+
+	reader, err := db.Topic().StartListener("consumer", handler, topicoptions.ReadTopic("my-topic"))
+	if err != nil {
+		return nil, err
+	}
+	if err = reader.WaitInit(ctx); err != nil {
+		return nil, err
+	}
+
+	return reader, nil
+}
+
+type TopicEventsHandler struct {
+	topiclistener.BaseHandler
+	listener *topiclistener.TopicListener
+
+	m     sync.Mutex
+	locks map[int64]int64 // [partitionSessionID]lockID
+}
+
+func (h *TopicEventsHandler) OnReaderCreated(req *topiclistener.ReaderReady) error {
+	h.listener = req.Listener
+
+	return nil
+}
+
+func (h *TopicEventsHandler) OnReadMessages(
+	ctx context.Context,
+	event *topiclistener.ReadMessages,
+) error {
+	for _, mess := range event.Batch.Messages {
+		log.Printf("Receive message: %v/%v/%v", mess.Topic(), mess.PartitionID(), mess.SeqNo)
+	}
+
+	processBatch(ctx, event.Batch)
+
+	return nil
+}
+
+func (h *TopicEventsHandler) OnStartPartitionSessionRequest(
+	ctx context.Context,
+	event *topiclistener.StartPartitionSessionRequest,
+) error {
+	lockID, offset, err := lockPartition(ctx, event.PartitionSession.TopicPath, event.PartitionSession.PartitionID)
+
+	h.m.Lock()
+	h.locks[event.PartitionSession.PartitionSessionID] = lockID
+	h.m.Unlock()
+
+	log.Printf("Started read partition %v/%v", event.PartitionSession.TopicPath, event.PartitionSession.PartitionID)
+	event.ConfirmWithParams(
+		topiclistener.StartPartitionSessionConfirm{}.
+			WithReadOffet(offset).
+			WithCommitOffset(offset),
+	)
+
+	return err
+}
+
+func (h *TopicEventsHandler) OnStopPartitionSessionRequest(
+	ctx context.Context,
+	event *topiclistener.StopPartitionSessionRequest,
+) error {
+	h.m.Lock()
+	lockID := h.locks[event.PartitionSession.PartitionSessionID]
+	delete(h.locks, event.PartitionSession.PartitionSessionID)
+	h.m.Unlock()
+
+	err := unlockPartition(ctx, lockID)
+	event.Confirm()
+
+	return err
+}
+
+func lockPartition(ctx context.Context, topic string, partitionID int64) (lockID, offset int64, err error) {
+	panic("not implemented in the example")
+}
+
+func unlockPartition(ctx context.Context, lockID int64) error {
+	panic("not implemented in the example")
+}

internal/background/worker.go

@@ -122,6 +122,10 @@ func (b *Worker) CloseReason() error {
 	return b.closeReason
 }
 
+func (b *Worker) StopDone() <-chan empty.Struct {
+	return b.tasksCompleted
+}
+
 func (b *Worker) init() {
 	b.onceInit.Do(func() {
 		if b.ctx == nil {

internal/grpcwrapper/rawtopic/rawtopicreader/rawtopicreader_test.go

@@ -0,0 +1,4 @@
+package rawtopicreader
+
+// Check interface implementation
+var _ TopicReaderStreamInterface = StreamReader{}

internal/grpcwrapper/rawtopic/rawtopicreader/stream_reader_stream_interface.go

@@ -0,0 +1,9 @@
+package rawtopicreader
+
+//go:generate mockgen -destination ../rawtopicreadermock/stream_reader_stream_interface_mock.go -package rawtopicreadermock -write_package_comment=false --typed . TopicReaderStreamInterface
+
+type TopicReaderStreamInterface interface {
+	Recv() (ServerMessage, error)
+	Send(msg ClientMessage) error
+	CloseSend() error
+}

internal/grpcwrapper/rawtopic/rawtopicreadermock/stream_reader_stream_interface_mock.go

@@ -10,9 +10,12 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawydb"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topiclistenerinternal"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicreadercommon"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicreaderinternal"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicwriterinternal"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
+	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topiclistener"
 	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicoptions"
 	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicreader"
 	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topictypes"
@@ -202,6 +205,38 @@ func (c *Client) Drop(ctx context.Context, path string, opts ...topicoptions.Dro
 	return call(ctx)
 }
 
+// StartListener starts read listen topic with the handler
+// it is fast non block call, connection starts in background
+func (c *Client) StartListener(
+	consumer string,
+	handler topiclistener.EventHandler,
+	readSelectors topicoptions.ReadSelectors,
+	opts ...topicoptions.ListenerOption,
+) (*topiclistener.TopicListener, error) {
+	cfg := topiclistenerinternal.NewStreamListenerConfig()
+
+	cfg.Consumer = consumer
+
+	cfg.Selectors = make([]*topicreadercommon.PublicReadSelector, len(readSelectors))
+	for i := range readSelectors {
+		cfg.Selectors[i] = readSelectors[i].Clone()
+	}
+
+	for _, opt := range opts {
+		if opt == nil {
+			continue
+		}
+
+		opt(&cfg)
+	}
+
+	if err := cfg.Validate(); err != nil {
+		return nil, err
+	}
+
+	return topiclistener.NewTopicListener(&c.rawClient, &cfg, handler)
+}
+
 // StartReader create new topic reader and start pull messages from server
 // it is fast non block call, connection will start in background
 func (c *Client) StartReader(