aboutsummaryrefslogtreecommitdiff
path: root/weed/mq/client/sub_client/subscriber.go
blob: 68bf74c5eb2a82c88e38214c6b9ff3d3a6b0029d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
package sub_client

import (
	"context"
	"sync"

	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
	"google.golang.org/grpc"
)

type SubscriberConfiguration struct {
	ClientId                string
	ConsumerGroup           string
	ConsumerGroupInstanceId string
	GrpcDialOption          grpc.DialOption
	MaxPartitionCount       int32 // how many partitions to process concurrently
	SlidingWindowSize       int32 // how many messages to process concurrently per partition
}

func (s *SubscriberConfiguration) String() string {
	return "ClientId: " + s.ClientId + ", ConsumerGroup: " + s.ConsumerGroup + ", ConsumerGroupInstanceId: " + s.ConsumerGroupInstanceId
}

type ContentConfiguration struct {
	Topic            topic.Topic
	Filter           string
	PartitionOffsets []*schema_pb.PartitionOffset
	OffsetType       schema_pb.OffsetType
	OffsetTsNs       int64
}

type OnDataMessageFn func(m *mq_pb.SubscribeMessageResponse_Data)
type OnCompletionFunc func()

type TopicSubscriber struct {
	ctx                              context.Context
	SubscriberConfig                 *SubscriberConfiguration
	ContentConfig                    *ContentConfiguration
	brokerPartitionAssignmentChan    chan *mq_pb.SubscriberToSubCoordinatorResponse
	brokerPartitionAssignmentAckChan chan *mq_pb.SubscriberToSubCoordinatorRequest
	OnDataMessageFunc                OnDataMessageFn
	OnCompletionFunc                 OnCompletionFunc
	bootstrapBrokers                 []string
	activeProcessors                 map[topic.Partition]*ProcessorState
	activeProcessorsLock             sync.Mutex
	PartitionOffsetChan              chan KeyedTimestamp
}

func NewTopicSubscriber(ctx context.Context, bootstrapBrokers []string, subscriber *SubscriberConfiguration, content *ContentConfiguration, partitionOffsetChan chan KeyedTimestamp) *TopicSubscriber {
	return &TopicSubscriber{
		ctx:                              ctx,
		SubscriberConfig:                 subscriber,
		ContentConfig:                    content,
		brokerPartitionAssignmentChan:    make(chan *mq_pb.SubscriberToSubCoordinatorResponse, 1024),
		brokerPartitionAssignmentAckChan: make(chan *mq_pb.SubscriberToSubCoordinatorRequest, 1024),
		bootstrapBrokers:                 bootstrapBrokers,
		activeProcessors:                 make(map[topic.Partition]*ProcessorState),
		PartitionOffsetChan:              partitionOffsetChan,
	}
}

func (sub *TopicSubscriber) SetOnDataMessageFn(fn OnDataMessageFn) {
	sub.OnDataMessageFunc = fn
}

func (sub *TopicSubscriber) SetCompletionFunc(onCompletionFn OnCompletionFunc) {
	sub.OnCompletionFunc = onCompletionFn
}