Accumulated changes for message queue (#6600)

* rename

* set agent address

* refactor

* add agent sub

* pub messages

* grpc new client

* can publish records via agent

* send init message with session id

* fmt

* check cancelled request while waiting

* use sessionId

* handle possible nil stream

* subscriber process messages

* separate debug port

* use atomic int64

* less logs

* minor

* skip io.EOF

* rename

* remove unused

* use saved offsets

* do not reuse session, since always session id is new after restart

remove last active ts from SessionEntry

* simplify printing

* purge unused

* just proxy the subscription, skipping the session step

* adjust offset types

* subscribe offset type and possible value

* start after the known tsns

* avoid wrongly set startPosition

* move

* remove

* refactor

* typo

* fix

* fix changed path

weed/mq/agent/agent_grpc_pub_session.go

@@ -7,13 +7,12 @@ import (
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_agent_pb"
 	"log/slog"
 	"math/rand/v2"
-	"time"
 )
 
 func (a *MessageQueueAgent) StartPublishSession(ctx context.Context, req *mq_agent_pb.StartPublishSessionRequest) (*mq_agent_pb.StartPublishSessionResponse, error) {
 	sessionId := rand.Int64()
 
-	topicPublisher := pub_client.NewTopicPublisher(
+	topicPublisher, err := pub_client.NewTopicPublisher(
 		&pub_client.PublisherConfiguration{
 			Topic:          topic.NewTopic(req.Topic.Namespace, req.Topic.Name),
 			PartitionCount: req.PartitionCount,
@@ -21,18 +20,11 @@ func (a *MessageQueueAgent) StartPublishSession(ctx context.Context, req *mq_age
 			PublisherName:  req.PublisherName,
 			RecordType:     req.RecordType,
 		})
+	if err != nil {
+		return nil, err
+	}
 
 	a.publishersLock.Lock()
-	// remove inactive publishers to avoid memory leak
-	for k, entry := range a.publishers {
-		if entry.lastActiveTsNs == 0 {
-			// this is an active session
-			continue
-		}
-		if time.Unix(0, entry.lastActiveTsNs).Add(10 * time.Hour).Before(time.Now()) {
-			delete(a.publishers, k)
-		}
-	}
 	a.publishers[SessionId(sessionId)] = &SessionEntry[*pub_client.TopicPublisher]{
 		entry: topicPublisher,
 	}

weed/mq/agent/agent_grpc_publish.go

@@ -3,24 +3,25 @@ package agent
 import (
 	"fmt"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_agent_pb"
-	"time"
 )
 
-func (a *MessageQueueAgent) PublishRecordRequest(stream mq_agent_pb.SeaweedMessagingAgent_PublishRecordServer) error {
+func (a *MessageQueueAgent) PublishRecord(stream mq_agent_pb.SeaweedMessagingAgent_PublishRecordServer) error {
 	m, err := stream.Recv()
 	if err != nil {
 		return err
 	}
+	sessionId := SessionId(m.SessionId)
 	a.publishersLock.RLock()
-	publisherEntry, found := a.publishers[SessionId(m.SessionId)]
+	publisherEntry, found := a.publishers[sessionId]
 	a.publishersLock.RUnlock()
 	if !found {
-		return fmt.Errorf("publish session id %d not found", m.SessionId)
+		return fmt.Errorf("publish session id %d not found", sessionId)
 	}
 	defer func() {
-		publisherEntry.lastActiveTsNs = time.Now().UnixNano()
+		a.publishersLock.Lock()
+		delete(a.publishers, sessionId)
+		a.publishersLock.Unlock()
 	}()
-	publisherEntry.lastActiveTsNs = 0
 
 	if m.Value != nil {
 		if err := publisherEntry.entry.PublishRecord(m.Key, m.Value); err != nil {
@@ -29,7 +30,7 @@ func (a *MessageQueueAgent) PublishRecordRequest(stream mq_agent_pb.SeaweedMessa
 	}
 
 	for {
-		m, err := stream.Recv()
+		m, err = stream.Recv()
 		if err != nil {
 			return err
 		}

weed/mq/agent/agent_grpc_sub_session.go

@@ -1,57 +0,0 @@
-package agent
-
-import (
-	"context"
-	"github.com/seaweedfs/seaweedfs/weed/mq/client/sub_client"
-	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
-	"github.com/seaweedfs/seaweedfs/weed/pb/mq_agent_pb"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials/insecure"
-	"math/rand/v2"
-	"time"
-)
-
-func (a *MessageQueueAgent) StartSubscribeSession(ctx context.Context, req *mq_agent_pb.StartSubscribeSessionRequest) (*mq_agent_pb.StartSubscribeSessionResponse, error) {
-	sessionId := rand.Int64()
-
-	subscriberConfig := &sub_client.SubscriberConfiguration{
-		ConsumerGroup:           req.ConsumerGroup,
-		ConsumerGroupInstanceId: req.ConsumerGroupInstanceId,
-		GrpcDialOption:          grpc.WithTransportCredentials(insecure.NewCredentials()),
-		MaxPartitionCount:       req.MaxSubscribedPartitions,
-		SlidingWindowSize:       req.SlidingWindowSize,
-	}
-
-	contentConfig := &sub_client.ContentConfiguration{
-		Topic:            topic.FromPbTopic(req.Topic),
-		Filter:           req.Filter,
-		PartitionOffsets: req.PartitionOffsets,
-	}
-
-	topicSubscriber := sub_client.NewTopicSubscriber(
-		a.brokersList(),
-		subscriberConfig,
-		contentConfig,
-		make(chan sub_client.KeyedOffset, 1024),
-	)
-
-	a.subscribersLock.Lock()
-	// remove inactive publishers to avoid memory leak
-	for k, entry := range a.subscribers {
-		if entry.lastActiveTsNs == 0 {
-			// this is an active session
-			continue
-		}
-		if time.Unix(0, entry.lastActiveTsNs).Add(10 * time.Hour).Before(time.Now()) {
-			delete(a.subscribers, k)
-		}
-	}
-	a.subscribers[SessionId(sessionId)] = &SessionEntry[*sub_client.TopicSubscriber]{
-		entry: topicSubscriber,
-	}
-	a.subscribersLock.Unlock()
-
-	return &mq_agent_pb.StartSubscribeSessionResponse{
-		SessionId: sessionId,
-	}, nil
-}

weed/mq/agent/agent_grpc_subscribe.go

@@ -1,59 +1,59 @@
 package agent
 
 import (
-	"fmt"
+	"context"
 	"github.com/seaweedfs/seaweedfs/weed/glog"
 	"github.com/seaweedfs/seaweedfs/weed/mq/client/sub_client"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_agent_pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+	"github.com/seaweedfs/seaweedfs/weed/util"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/protobuf/proto"
-	"time"
 )
 
-func (a *MessageQueueAgent) SubscribeRecordRequest(stream mq_agent_pb.SeaweedMessagingAgent_SubscribeRecordServer) error {
+func (a *MessageQueueAgent) SubscribeRecord(stream mq_agent_pb.SeaweedMessagingAgent_SubscribeRecordServer) error {
 	// the first message is the subscribe request
 	// it should only contain the session id
-	m, err := stream.Recv()
+	initMessage, err := stream.Recv()
 	if err != nil {
 		return err
 	}
-	a.subscribersLock.RLock()
-	subscriberEntry, found := a.subscribers[SessionId(m.SessionId)]
-	a.subscribersLock.RUnlock()
-	if !found {
-		return fmt.Errorf("subscribe session id %d not found", m.SessionId)
-	}
-	defer func() {
-		subscriberEntry.lastActiveTsNs = time.Now().UnixNano()
-	}()
-	subscriberEntry.lastActiveTsNs = 0
+
+	subscriber := a.handleInitSubscribeRecordRequest(stream.Context(), initMessage.Init)
 
 	var lastErr error
-	subscriberEntry.entry.SetOnDataMessageFn(func(m *mq_pb.SubscribeMessageResponse_Data) {
-		record := &schema_pb.RecordValue{}
-		err := proto.Unmarshal(m.Data.Value, record)
-		if err != nil {
-			if lastErr == nil {
-				lastErr = err
+	executors := util.NewLimitedConcurrentExecutor(int(subscriber.SubscriberConfig.SlidingWindowSize))
+	subscriber.SetOnDataMessageFn(func(m *mq_pb.SubscribeMessageResponse_Data) {
+		executors.Execute(func() {
+			record := &schema_pb.RecordValue{}
+			err := proto.Unmarshal(m.Data.Value, record)
+			if err != nil {
+				glog.V(0).Infof("unmarshal record value: %v", err)
+				if lastErr == nil {
+					lastErr = err
+				}
+				return
 			}
-			return
-		}
-		if sendErr := stream.Send(&mq_agent_pb.SubscribeRecordResponse{
-			Key:   m.Data.Key,
-			Value: record,
-			TsNs:  m.Data.TsNs,
-		}); sendErr != nil {
-			if lastErr == nil {
-				lastErr = sendErr
+			if sendErr := stream.Send(&mq_agent_pb.SubscribeRecordResponse{
+				Key:   m.Data.Key,
+				Value: record,
+				TsNs:  m.Data.TsNs,
+			}); sendErr != nil {
+				glog.V(0).Infof("send record: %v", sendErr)
+				if lastErr == nil {
+					lastErr = sendErr
+				}
 			}
-		}
+		})
 	})
 
 	go func() {
-		subErr := subscriberEntry.entry.Subscribe()
+		subErr := subscriber.Subscribe()
 		if subErr != nil {
-			glog.V(0).Infof("subscriber %d subscribe: %v", m.SessionId, subErr)
+			glog.V(0).Infof("subscriber %s subscribe: %v", subscriber.SubscriberConfig.String(), subErr)
 			if lastErr == nil {
 				lastErr = subErr
 			}
@@ -63,13 +63,43 @@ func (a *MessageQueueAgent) SubscribeRecordRequest(stream mq_agent_pb.SeaweedMes
 	for {
 		m, err := stream.Recv()
 		if err != nil {
+			glog.V(0).Infof("subscriber %s receive: %v", subscriber.SubscriberConfig.String(), err)
 			return err
 		}
 		if m != nil {
-			subscriberEntry.entry.PartitionOffsetChan <- sub_client.KeyedOffset{
+			subscriber.PartitionOffsetChan <- sub_client.KeyedOffset{
 				Key:    m.AckKey,
 				Offset: m.AckSequence,
 			}
 		}
 	}
 }
+
+func (a *MessageQueueAgent) handleInitSubscribeRecordRequest(ctx context.Context, req *mq_agent_pb.SubscribeRecordRequest_InitSubscribeRecordRequest) *sub_client.TopicSubscriber {
+
+	subscriberConfig := &sub_client.SubscriberConfiguration{
+		ConsumerGroup:           req.ConsumerGroup,
+		ConsumerGroupInstanceId: req.ConsumerGroupInstanceId,
+		GrpcDialOption:          grpc.WithTransportCredentials(insecure.NewCredentials()),
+		MaxPartitionCount:       req.MaxSubscribedPartitions,
+		SlidingWindowSize:       req.SlidingWindowSize,
+	}
+
+	contentConfig := &sub_client.ContentConfiguration{
+		Topic:            topic.FromPbTopic(req.Topic),
+		Filter:           req.Filter,
+		PartitionOffsets: req.PartitionOffsets,
+		OffsetType:       req.OffsetType,
+		OffsetTsNs:       req.OffsetTsNs,
+	}
+
+	topicSubscriber := sub_client.NewTopicSubscriber(
+		ctx,
+		a.brokersList(),
+		subscriberConfig,
+		contentConfig,
+		make(chan sub_client.KeyedOffset, 1024),
+	)
+
+	return topicSubscriber
+}

weed/mq/agent/agent_server.go

@@ -11,8 +11,7 @@ import (
 
 type SessionId int64
 type SessionEntry[T any] struct {
-	entry          T
-	lastActiveTsNs int64
+	entry T
 }
 
 type MessageQueueAgentOptions struct {
@@ -32,11 +31,15 @@ type MessageQueueAgent struct {
 
 func NewMessageQueueAgent(option *MessageQueueAgentOptions, grpcDialOption grpc.DialOption) *MessageQueueAgent {
 
-	// check masters to list all brokers
+	// initialize brokers which may change later
+	var brokers []pb.ServerAddress
+	for _, broker := range option.SeedBrokers {
+		brokers = append(brokers, broker)
+	}
 
 	return &MessageQueueAgent{
 		option:         option,
-		brokers:        []pb.ServerAddress{},
+		brokers:        brokers,
 		grpcDialOption: grpcDialOption,
 		publishers:     make(map[SessionId]*SessionEntry[*pub_client.TopicPublisher]),
 		subscribers:    make(map[SessionId]*SessionEntry[*sub_client.TopicSubscriber]),

weed/mq/broker/broker_grpc_sub.go

@@ -68,7 +68,9 @@ func (b *MessageQueueBroker) SubscribeMessage(stream mq_pb.SeaweedMessaging_Subs
 		} else {
 			defer func() {
 				println("closing SubscribeFollowMe connection", follower)
-				subscribeFollowMeStream.CloseSend()
+				if subscribeFollowMeStream != nil {
+					subscribeFollowMeStream.CloseSend()
+				}
 				// followerGrpcConnection.Close()
 			}()
 			followerClient := mq_pb.NewSeaweedMessagingClient(followerGrpcConnection)
@@ -142,7 +144,9 @@ func (b *MessageQueueBroker) SubscribeMessage(stream mq_pb.SeaweedMessaging_Subs
 					Close: &mq_pb.SubscribeFollowMeRequest_CloseMessage{},
 				},
 			}); err != nil {
-				glog.Errorf("Error sending close to follower: %v", err)
+				if err != io.EOF {
+					glog.Errorf("Error sending close to follower: %v", err)
+				}
 			}
 		}
 	}()
@@ -178,6 +182,19 @@ func (b *MessageQueueBroker) SubscribeMessage(stream mq_pb.SeaweedMessaging_Subs
 
 		for imt.IsInflight(logEntry.Key) {
 			time.Sleep(137 * time.Millisecond)
+			// Check if the client has disconnected by monitoring the context
+			select {
+			case <-ctx.Done():
+				err := ctx.Err()
+				if err == context.Canceled {
+					// Client disconnected
+					return false, nil
+				}
+				glog.V(0).Infof("Subscriber %s disconnected: %v", clientName, err)
+				return false, nil
+			default:
+				// Continue processing the request
+			}
 		}
 		if logEntry.Key != nil {
 			imt.EnflightMessage(logEntry.Key, logEntry.TsNs)
@@ -204,20 +221,35 @@ func (b *MessageQueueBroker) getRequestPosition(initMessage *mq_pb.SubscribeMess
 		return
 	}
 	offset := initMessage.GetPartitionOffset()
-	if offset.StartTsNs != 0 {
+	offsetType := initMessage.OffsetType
+
+	// reset to earliest or latest
+	if offsetType == schema_pb.OffsetType_RESET_TO_EARLIEST {
+		startPosition = log_buffer.NewMessagePosition(1, -3)
+		return
+	}
+	if offsetType == schema_pb.OffsetType_RESET_TO_LATEST {
+		startPosition = log_buffer.NewMessagePosition(time.Now().UnixNano(), -4)
+		return
+	}
+
+	// use the exact timestamp
+	if offsetType == schema_pb.OffsetType_EXACT_TS_NS {
 		startPosition = log_buffer.NewMessagePosition(offset.StartTsNs, -2)
 		return
 	}
+
+	// try to resume
 	if storedOffset, err := b.readConsumerGroupOffset(initMessage); err == nil {
 		glog.V(0).Infof("resume from saved offset %v %v %v: %v", initMessage.Topic, initMessage.PartitionOffset.Partition, initMessage.ConsumerGroup, storedOffset)
 		startPosition = log_buffer.NewMessagePosition(storedOffset, -2)
 		return
 	}
 
-	if offset.StartType == schema_pb.PartitionOffsetStartType_EARLIEST {
-		startPosition = log_buffer.NewMessagePosition(1, -3)
-	} else if offset.StartType == schema_pb.PartitionOffsetStartType_LATEST {
-		startPosition = log_buffer.NewMessagePosition(time.Now().UnixNano(), -4)
+	if offsetType == schema_pb.OffsetType_RESUME_OR_EARLIEST {
+		startPosition = log_buffer.NewMessagePosition(1, -5)
+	} else if offsetType == schema_pb.OffsetType_RESUME_OR_LATEST {
+		startPosition = log_buffer.NewMessagePosition(time.Now().UnixNano(), -6)
 	}
 	return
 }

weed/mq/client/agent_client/agent_publish.go

@@ -1,14 +0,0 @@
-package agent_client
-
-import (
-	"github.com/seaweedfs/seaweedfs/weed/pb/mq_agent_pb"
-	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
-)
-
-func (a *PublishSession) PublishMessageRecord(key []byte, record *schema_pb.RecordValue) error {
-	return a.stream.Send(&mq_agent_pb.PublishRecordRequest{
-		SessionId: a.sessionId,
-		Key:       key,
-		Value:     record,
-	})
-}

weed/mq/client/agent_client/agent_subscribe.go

@@ -1,17 +0,0 @@
-package agent_client
-
-import (
-	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
-)
-
-func (a *SubscribeSession) SubscribeMessageRecord(
-	onEachMessageFn func(key []byte, record *schema_pb.RecordValue),
-	onCompletionFn func()) error {
-	for {
-		resp, err := a.stream.Recv()
-		if err != nil {
-			return err
-		}
-		onEachMessageFn(resp.Key, resp.Value)
-	}
-}

weed/mq/client/agent_client/publish_session.go

@@ -4,10 +4,10 @@ import (
 	"context"
 	"fmt"
 	"github.com/seaweedfs/seaweedfs/weed/mq/schema"
-	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_agent_pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
 )
 
 type PublishSession struct {
@@ -15,13 +15,12 @@ type PublishSession struct {
 	partitionCount int
 	publisherName  string
 	stream         grpc.BidiStreamingClient[mq_agent_pb.PublishRecordRequest, mq_agent_pb.PublishRecordResponse]
-	sessionId      int64
 }
 
 func NewPublishSession(agentAddress string, topicSchema *schema.Schema, partitionCount int, publisherName string) (*PublishSession, error) {
 
 	// call local agent grpc server to create a new session
-	clientConn, err := pb.GrpcDial(context.Background(), agentAddress, true, grpc.WithInsecure())
+	clientConn, err := grpc.NewClient(agentAddress, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
 		return nil, fmt.Errorf("dial agent server %s: %v", agentAddress, err)
 	}
@@ -48,12 +47,17 @@ func NewPublishSession(agentAddress string, topicSchema *schema.Schema, partitio
 		return nil, fmt.Errorf("publish record: %v", err)
 	}
 
+	if err = stream.Send(&mq_agent_pb.PublishRecordRequest{
+		SessionId: resp.SessionId,
+	}); err != nil {
+		return nil, fmt.Errorf("send session id: %v", err)
+	}
+
 	return &PublishSession{
 		schema:         topicSchema,
 		partitionCount: partitionCount,
 		publisherName:  publisherName,
 		stream:         stream,
-		sessionId:      resp.SessionId,
 	}, nil
 }
 
@@ -68,3 +72,10 @@ func (a *PublishSession) CloseSession() error {
 	a.schema = nil
 	return err
 }
+
+func (a *PublishSession) PublishMessageRecord(key []byte, record *schema_pb.RecordValue) error {
+	return a.stream.Send(&mq_agent_pb.PublishRecordRequest{
+		Key:   key,
+		Value: record,
+	})
+}

weed/mq/client/agent_client/subscribe_session.go

@@ -4,50 +4,48 @@ import (
 	"context"
 	"fmt"
 	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
-	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_agent_pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
 )
 
 type SubscribeOption struct {
 	ConsumerGroup           string
 	ConsumerGroupInstanceId string
 	Topic                   topic.Topic
+	OffsetType              schema_pb.OffsetType
+	OffsetTsNs              int64
 	Filter                  string
 	MaxSubscribedPartitions int32
-	PerPartitionConcurrency int32
+	SlidingWindowSize       int32
 }
 
 type SubscribeSession struct {
-	Option    *SubscribeOption
-	stream    grpc.BidiStreamingClient[mq_agent_pb.SubscribeRecordRequest, mq_agent_pb.SubscribeRecordResponse]
-	sessionId int64
+	Option *SubscribeOption
+	stream grpc.BidiStreamingClient[mq_agent_pb.SubscribeRecordRequest, mq_agent_pb.SubscribeRecordResponse]
 }
 
 func NewSubscribeSession(agentAddress string, option *SubscribeOption) (*SubscribeSession, error) {
 	// call local agent grpc server to create a new session
-	clientConn, err := pb.GrpcDial(context.Background(), agentAddress, true, grpc.WithInsecure())
+	clientConn, err := grpc.NewClient(agentAddress, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
 		return nil, fmt.Errorf("dial agent server %s: %v", agentAddress, err)
 	}
 	agentClient := mq_agent_pb.NewSeaweedMessagingAgentClient(clientConn)
 
-	resp, err := agentClient.StartSubscribeSession(context.Background(), &mq_agent_pb.StartSubscribeSessionRequest{
+	initRequest := &mq_agent_pb.SubscribeRecordRequest_InitSubscribeRecordRequest{
 		ConsumerGroup:           option.ConsumerGroup,
 		ConsumerGroupInstanceId: option.ConsumerGroupInstanceId,
 		Topic: &schema_pb.Topic{
 			Namespace: option.Topic.Namespace,
 			Name:      option.Topic.Name,
 		},
+		OffsetType:              option.OffsetType,
+		OffsetTsNs:              option.OffsetTsNs,
 		MaxSubscribedPartitions: option.MaxSubscribedPartitions,
 		Filter:                  option.Filter,
-	})
-	if err != nil {
-		return nil, err
-	}
-	if resp.Error != "" {
-		return nil, fmt.Errorf("start subscribe session: %v", resp.Error)
+		SlidingWindowSize:       option.SlidingWindowSize,
 	}
 
 	stream, err := agentClient.SubscribeRecord(context.Background())
@@ -55,9 +53,35 @@ func NewSubscribeSession(agentAddress string, option *SubscribeOption) (*Subscri
 		return nil, fmt.Errorf("subscribe record: %v", err)
 	}
 
+	if err = stream.Send(&mq_agent_pb.SubscribeRecordRequest{
+		Init: initRequest,
+	}); err != nil {
+		return nil, fmt.Errorf("send session id: %v", err)
+	}
+
 	return &SubscribeSession{
-		Option:    option,
-		stream:    stream,
-		sessionId: resp.SessionId,
+		Option: option,
+		stream: stream,
 	}, nil
 }
+
+func (s *SubscribeSession) CloseSession() error {
+	err := s.stream.CloseSend()
+	return err
+}
+
+func (a *SubscribeSession) SubscribeMessageRecord(
+	onEachMessageFn func(key []byte, record *schema_pb.RecordValue),
+	onCompletionFn func()) error {
+	for {
+		resp, err := a.stream.Recv()
+		if err != nil {
+			return err
+		}
+		onEachMessageFn(resp.Key, resp.Value)
+	}
+	if onCompletionFn != nil {
+		onCompletionFn()
+	}
+	return nil
+}

weed/mq/client/cmd/agent_pub_record/publisher_record.go

@@ -1,128 +0,0 @@
-package main
-
-import (
-	"flag"
-	"fmt"
-	"github.com/seaweedfs/seaweedfs/weed/mq/client/agent_client"
-	"github.com/seaweedfs/seaweedfs/weed/mq/schema"
-	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
-	"log"
-	"sync"
-	"sync/atomic"
-	"time"
-)
-
-var (
-	messageCount   = flag.Int("n", 1000, "message count")
-	messageDelay   = flag.Duration("d", time.Second, "delay between messages")
-	concurrency    = flag.Int("c", 4, "concurrent publishers")
-	partitionCount = flag.Int("p", 6, "partition count")
-
-	clientName = flag.String("client", "c1", "client name")
-
-	namespace   = flag.String("ns", "test", "namespace")
-	t           = flag.String("t", "test", "t")
-	seedBrokers = flag.String("brokers", "localhost:17777", "seed brokers")
-
-	counter int32
-)
-
-func doPublish(publisher *agent_client.PublishSession, id int) {
-	startTime := time.Now()
-	for {
-		i := atomic.AddInt32(&counter, 1)
-		if i > int32(*messageCount) {
-			break
-		}
-		// Simulate publishing a message
-		myRecord := genMyRecord(int32(i))
-		if err := publisher.PublishMessageRecord(myRecord.Key, myRecord.ToRecordValue()); err != nil {
-			fmt.Println(err)
-			break
-		}
-		if *messageDelay > 0 {
-			time.Sleep(*messageDelay)
-			fmt.Printf("sent %+v\n", string(myRecord.Key))
-		}
-	}
-	elapsed := time.Since(startTime)
-	log.Printf("Publisher %s-%d finished in %s", *clientName, id, elapsed)
-}
-
-type MyRecord struct {
-	Key    []byte
-	Field1 []byte
-	Field2 string
-	Field3 int32
-	Field4 int64
-	Field5 float32
-	Field6 float64
-	Field7 bool
-}
-
-func genMyRecord(id int32) *MyRecord {
-	return &MyRecord{
-		Key:    []byte(fmt.Sprintf("key-%s-%d", *clientName, id)),
-		Field1: []byte(fmt.Sprintf("field1-%s-%d", *clientName, id)),
-		Field2: fmt.Sprintf("field2-%s-%d", *clientName, id),
-		Field3: id,
-		Field4: int64(id),
-		Field5: float32(id),
-		Field6: float64(id),
-		Field7: id%2 == 0,
-	}
-}
-
-func (r *MyRecord) ToRecordValue() *schema_pb.RecordValue {
-	return schema.RecordBegin().
-		SetBytes("key", r.Key).
-		SetBytes("field1", r.Field1).
-		SetString("field2", r.Field2).
-		SetInt32("field3", r.Field3).
-		SetInt64("field4", r.Field4).
-		SetFloat("field5", r.Field5).
-		SetDouble("field6", r.Field6).
-		SetBool("field7", r.Field7).
-		RecordEnd()
-}
-
-func main() {
-	flag.Parse()
-
-	recordType := schema.RecordTypeBegin().
-		WithField("key", schema.TypeBytes).
-		WithField("field1", schema.TypeBytes).
-		WithField("field2", schema.TypeString).
-		WithField("field3", schema.TypeInt32).
-		WithField("field4", schema.TypeInt64).
-		WithField("field5", schema.TypeFloat).
-		WithField("field6", schema.TypeDouble).
-		WithField("field7", schema.TypeBoolean).
-		RecordTypeEnd()
-
-	session, err := agent_client.NewPublishSession("localhost:16777", schema.NewSchema(*namespace, *t, recordType), *partitionCount, *clientName)
-	if err != nil {
-		log.Printf("failed to create session: %v", err)
-		return
-	}
-	defer session.CloseSession()
-
-	startTime := time.Now()
-
-	var wg sync.WaitGroup
-	// Start multiple publishers
-	for i := 0; i < *concurrency; i++ {
-		wg.Add(1)
-		go func(id int) {
-			defer wg.Done()
-			doPublish(session, id)
-		}(i)
-	}
-
-	// Wait for all publishers to finish
-	wg.Wait()
-	elapsed := time.Since(startTime)
-
-	log.Printf("Published %d messages in %s (%.2f msg/s)", *messageCount, elapsed, float64(*messageCount)/elapsed.Seconds())
-
-}

weed/mq/client/cmd/weed_pub_kv/publisher_kv.go

@@ -1,78 +0,0 @@
-package main
-
-import (
-	"flag"
-	"fmt"
-	"github.com/seaweedfs/seaweedfs/weed/mq/client/pub_client"
-	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
-	util_http "github.com/seaweedfs/seaweedfs/weed/util/http"
-	"log"
-	"strings"
-	"sync"
-	"time"
-)
-
-var (
-	messageCount   = flag.Int("n", 1000, "message count")
-	concurrency    = flag.Int("c", 4, "concurrent publishers")
-	partitionCount = flag.Int("p", 6, "partition count")
-
-	clientName = flag.String("client", "c1", "client name")
-
-	namespace   = flag.String("ns", "test", "namespace")
-	t           = flag.String("t", "test", "t")
-	seedBrokers = flag.String("brokers", "localhost:17777", "seed brokers")
-)
-
-func doPublish(publisher *pub_client.TopicPublisher, id int) {
-	startTime := time.Now()
-	for i := 0; i < *messageCount / *concurrency; i++ {
-		// Simulate publishing a message
-		key := []byte(fmt.Sprintf("key-%s-%d-%d", *clientName, id, i))
-		value := []byte(fmt.Sprintf("value-%s-%d-%d", *clientName, id, i))
-		if err := publisher.Publish(key, value); err != nil {
-			fmt.Println(err)
-			break
-		}
-		time.Sleep(time.Second)
-		// println("Published", string(key), string(value))
-	}
-	if err := publisher.FinishPublish(); err != nil {
-		fmt.Println(err)
-	}
-	elapsed := time.Since(startTime)
-	log.Printf("Publisher %s-%d finished in %s", *clientName, id, elapsed)
-}
-
-func main() {
-	flag.Parse()
-	util_http.InitGlobalHttpClient()
-
-	config := &pub_client.PublisherConfiguration{
-		Topic:          topic.NewTopic(*namespace, *t),
-		PartitionCount: int32(*partitionCount),
-		Brokers:        strings.Split(*seedBrokers, ","),
-		PublisherName:  *clientName,
-	}
-	publisher := pub_client.NewTopicPublisher(config)
-
-	startTime := time.Now()
-
-	var wg sync.WaitGroup
-	// Start multiple publishers
-	for i := 0; i < *concurrency; i++ {
-		wg.Add(1)
-		go func(id int) {
-			defer wg.Done()
-			doPublish(publisher, id)
-		}(i)
-	}
-
-	// Wait for all publishers to finish
-	wg.Wait()
-	elapsed := time.Since(startTime)
-	publisher.Shutdown()
-
-	log.Printf("Published %d messages in %s (%.2f msg/s)", *messageCount, elapsed, float64(*messageCount)/elapsed.Seconds())
-
-}

weed/mq/client/cmd/weed_pub_record/publisher_record.go

@@ -1,138 +0,0 @@
-package main
-
-import (
-	"flag"
-	"fmt"
-	"github.com/seaweedfs/seaweedfs/weed/mq/client/pub_client"
-	"github.com/seaweedfs/seaweedfs/weed/mq/schema"
-	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
-	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
-	util_http "github.com/seaweedfs/seaweedfs/weed/util/http"
-	"log"
-	"strings"
-	"sync"
-	"sync/atomic"
-	"time"
-)
-
-var (
-	messageCount   = flag.Int("n", 1000, "message count")
-	messageDelay   = flag.Duration("d", time.Second, "delay between messages")
-	concurrency    = flag.Int("c", 4, "concurrent publishers")
-	partitionCount = flag.Int("p", 6, "partition count")
-
-	clientName = flag.String("client", "c1", "client name")
-
-	namespace   = flag.String("ns", "test", "namespace")
-	t           = flag.String("t", "test", "t")
-	seedBrokers = flag.String("brokers", "localhost:17777", "seed brokers")
-
-	counter int32
-)
-
-func doPublish(publisher *pub_client.TopicPublisher, id int) {
-	startTime := time.Now()
-	for {
-		i := atomic.AddInt32(&counter, 1)
-		if i > int32(*messageCount) {
-			break
-		}
-		// Simulate publishing a message
-		myRecord := genMyRecord(int32(i))
-		if err := publisher.PublishRecord(myRecord.Key, myRecord.ToRecordValue()); err != nil {
-			fmt.Println(err)
-			break
-		}
-		if *messageDelay > 0 {
-			time.Sleep(*messageDelay)
-			fmt.Printf("sent %+v\n", string(myRecord.Key))
-		}
-	}
-	if err := publisher.FinishPublish(); err != nil {
-		fmt.Println(err)
-	}
-	elapsed := time.Since(startTime)
-	log.Printf("Publisher %s-%d finished in %s", *clientName, id, elapsed)
-}
-
-type MyRecord struct {
-	Key    []byte
-	Field1 []byte
-	Field2 string
-	Field3 int32
-	Field4 int64
-	Field5 float32
-	Field6 float64
-	Field7 bool
-}
-
-func genMyRecord(id int32) *MyRecord {
-	return &MyRecord{
-		Key:    []byte(fmt.Sprintf("key-%s-%d", *clientName, id)),
-		Field1: []byte(fmt.Sprintf("field1-%s-%d", *clientName, id)),
-		Field2: fmt.Sprintf("field2-%s-%d", *clientName, id),
-		Field3: id,
-		Field4: int64(id),
-		Field5: float32(id),
-		Field6: float64(id),
-		Field7: id%2 == 0,
-	}
-}
-
-func (r *MyRecord) ToRecordValue() *schema_pb.RecordValue {
-	return schema.RecordBegin().
-		SetBytes("key", r.Key).
-		SetBytes("field1", r.Field1).
-		SetString("field2", r.Field2).
-		SetInt32("field3", r.Field3).
-		SetInt64("field4", r.Field4).
-		SetFloat("field5", r.Field5).
-		SetDouble("field6", r.Field6).
-		SetBool("field7", r.Field7).
-		RecordEnd()
-}
-
-func main() {
-	flag.Parse()
-	util_http.InitGlobalHttpClient()
-
-	recordType := schema.RecordTypeBegin().
-		WithField("key", schema.TypeBytes).
-		WithField("field1", schema.TypeBytes).
-		WithField("field2", schema.TypeString).
-		WithField("field3", schema.TypeInt32).
-		WithField("field4", schema.TypeInt64).
-		WithField("field5", schema.TypeFloat).
-		WithField("field6", schema.TypeDouble).
-		WithField("field7", schema.TypeBoolean).
-		RecordTypeEnd()
-
-	config := &pub_client.PublisherConfiguration{
-		Topic:          topic.NewTopic(*namespace, *t),
-		PartitionCount: int32(*partitionCount),
-		Brokers:        strings.Split(*seedBrokers, ","),
-		PublisherName:  *clientName,
-		RecordType:     recordType,
-	}
-	publisher := pub_client.NewTopicPublisher(config)
-
-	startTime := time.Now()
-
-	var wg sync.WaitGroup
-	// Start multiple publishers
-	for i := 0; i < *concurrency; i++ {
-		wg.Add(1)
-		go func(id int) {
-			defer wg.Done()
-			doPublish(publisher, id)
-		}(i)
-	}
-
-	// Wait for all publishers to finish
-	wg.Wait()
-	elapsed := time.Since(startTime)
-	publisher.Shutdown()
-
-	log.Printf("Published %d messages in %s (%.2f msg/s)", *messageCount, elapsed, float64(*messageCount)/elapsed.Seconds())
-
-}

weed/mq/client/cmd/weed_sub_kv/subscriber_kv.go

@@ -1,64 +0,0 @@
-package main
-
-import (
-	"flag"
-	"fmt"
-	"github.com/seaweedfs/seaweedfs/weed/glog"
-	"github.com/seaweedfs/seaweedfs/weed/mq/client/sub_client"
-	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
-	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
-	"github.com/seaweedfs/seaweedfs/weed/util"
-	util_http "github.com/seaweedfs/seaweedfs/weed/util/http"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials/insecure"
-	"strings"
-)
-
-var (
-	namespace               = flag.String("ns", "test", "namespace")
-	t                       = flag.String("topic", "test", "topic")
-	seedBrokers             = flag.String("brokers", "localhost:17777", "seed brokers")
-	maxPartitionCount       = flag.Int("maxPartitionCount", 3, "max partition count")
-	perPartitionConcurrency = flag.Int("perPartitionConcurrency", 1, "per partition concurrency")
-
-	clientId = flag.Uint("client_id", uint(util.RandomInt32()), "client id")
-)
-
-func main() {
-	flag.Parse()
-	util_http.InitGlobalHttpClient()
-
-	subscriberConfig := &sub_client.SubscriberConfiguration{
-		ConsumerGroup:           "test",
-		ConsumerGroupInstanceId: fmt.Sprintf("client-%d", *clientId),
-		GrpcDialOption:          grpc.WithTransportCredentials(insecure.NewCredentials()),
-		MaxPartitionCount:       int32(*maxPartitionCount),
-		SlidingWindowSize:       int32(*perPartitionConcurrency),
-	}
-
-	contentConfig := &sub_client.ContentConfiguration{
-		Topic:  topic.NewTopic(*namespace, *t),
-		Filter: "",
-	}
-
-	brokers := strings.Split(*seedBrokers, ",")
-	subscriber := sub_client.NewTopicSubscriber(brokers, subscriberConfig, contentConfig, make(chan sub_client.KeyedOffset, 1024))
-
-	counter := 0
-	executors := util.NewLimitedConcurrentExecutor(int(subscriberConfig.SlidingWindowSize))
-	subscriber.SetOnDataMessageFn(func(m *mq_pb.SubscribeMessageResponse_Data) {
-		executors.Execute(func() {
-			counter++
-			println(string(m.Data.Key), "=>", string(m.Data.Value), counter)
-		})
-	})
-
-	subscriber.SetCompletionFunc(func() {
-		glog.V(0).Infof("done received %d messages", counter)
-	})
-
-	if err := subscriber.Subscribe(); err != nil {
-		fmt.Println(err)
-	}
-
-}

weed/mq/client/cmd/weed_sub_record/subscriber_record.go

@@ -1,99 +0,0 @@
-package main
-
-import (
-	"flag"
-	"fmt"
-	"github.com/seaweedfs/seaweedfs/weed/glog"
-	"github.com/seaweedfs/seaweedfs/weed/mq/client/sub_client"
-	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
-	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
-	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
-	"github.com/seaweedfs/seaweedfs/weed/util"
-	util_http "github.com/seaweedfs/seaweedfs/weed/util/http"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials/insecure"
-	"google.golang.org/protobuf/proto"
-	"strings"
-	"time"
-)
-
-var (
-	namespace               = flag.String("ns", "test", "namespace")
-	t                       = flag.String("topic", "test", "topic")
-	seedBrokers             = flag.String("brokers", "localhost:17777", "seed brokers")
-	maxPartitionCount       = flag.Int("maxPartitionCount", 3, "max partition count")
-	perPartitionConcurrency = flag.Int("perPartitionConcurrency", 1, "per partition concurrency")
-	timeAgo                 = flag.Duration("timeAgo", 1*time.Hour, "start time before now. \"300ms\", \"1.5h\" or \"2h45m\". Valid time units are \"ns\", \"us\" (or \"µs\"), \"ms\", \"s\", \"m\", \"h\"")
-
-	clientId = flag.Uint("client_id", uint(util.RandomInt32()), "client id")
-)
-
-type MyRecord struct {
-	Key    []byte
-	Field1 []byte
-	Field2 string
-	Field3 int32
-	Field4 int64
-	Field5 float32
-	Field6 float64
-	Field7 bool
-}
-
-func FromSchemaRecordValue(recordValue *schema_pb.RecordValue) *MyRecord {
-	return &MyRecord{
-		Key:    recordValue.Fields["key"].GetBytesValue(),
-		Field1: recordValue.Fields["field1"].GetBytesValue(),
-		Field2: recordValue.Fields["field2"].GetStringValue(),
-		Field3: recordValue.Fields["field3"].GetInt32Value(),
-		Field4: recordValue.Fields["field4"].GetInt64Value(),
-		Field5: recordValue.Fields["field5"].GetFloatValue(),
-		Field6: recordValue.Fields["field6"].GetDoubleValue(),
-		Field7: recordValue.Fields["field7"].GetBoolValue(),
-	}
-}
-
-func main() {
-	flag.Parse()
-	util_http.InitGlobalHttpClient()
-
-	subscriberConfig := &sub_client.SubscriberConfiguration{
-		ConsumerGroup:           "test",
-		ConsumerGroupInstanceId: fmt.Sprintf("client-%d", *clientId),
-		GrpcDialOption:          grpc.WithTransportCredentials(insecure.NewCredentials()),
-		MaxPartitionCount:       int32(*maxPartitionCount),
-		SlidingWindowSize:       int32(*perPartitionConcurrency),
-	}
-
-	contentConfig := &sub_client.ContentConfiguration{
-		Topic:  topic.NewTopic(*namespace, *t),
-		Filter: "",
-		// StartTime: time.Now().Add(-*timeAgo),
-	}
-
-	brokers := strings.Split(*seedBrokers, ",")
-	subscriber := sub_client.NewTopicSubscriber(brokers, subscriberConfig, contentConfig, make(chan sub_client.KeyedOffset, 1024))
-
-	counter := 0
-	executors := util.NewLimitedConcurrentExecutor(int(subscriberConfig.SlidingWindowSize))
-	subscriber.SetOnDataMessageFn(func(m *mq_pb.SubscribeMessageResponse_Data) {
-		executors.Execute(func() {
-			counter++
-			record := &schema_pb.RecordValue{}
-			err := proto.Unmarshal(m.Data.Value, record)
-			if err != nil {
-				fmt.Printf("unmarshal record value: %v\n", err)
-			} else {
-				fmt.Printf("%s %d: %v\n", string(m.Data.Key), len(m.Data.Value), record)
-			}
-		})
-	})
-
-	subscriber.SetCompletionFunc(func() {
-		glog.V(0).Infof("done received %d messages", counter)
-	})
-
-	if err := subscriber.Subscribe(); err != nil {
-		fmt.Println(err)
-	}
-
-}

weed/mq/client/pub_client/publisher.go

@@ -34,8 +34,8 @@ type TopicPublisher struct {
 	jobs             []*EachPartitionPublishJob
 }
 
-func NewTopicPublisher(config *PublisherConfiguration) *TopicPublisher {
-	tp := &TopicPublisher{
+func NewTopicPublisher(config *PublisherConfiguration) (tp *TopicPublisher, err error) {
+	tp = &TopicPublisher{
 		partition2Buffer: interval.NewSearchTree[*buffered_queue.BufferedQueue[*mq_pb.DataMessage]](func(a, b int32) int {
 			return int(a - b)
 		}),
@@ -46,7 +46,7 @@ func NewTopicPublisher(config *PublisherConfiguration) *TopicPublisher {
 	wg := sync.WaitGroup{}
 	wg.Add(1)
 	go func() {
-		if err := tp.startSchedulerThread(&wg); err != nil {
+		if err = tp.startSchedulerThread(&wg); err != nil {
 			log.Println(err)
 			return
 		}
@@ -54,7 +54,7 @@ func NewTopicPublisher(config *PublisherConfiguration) *TopicPublisher {
 
 	wg.Wait()
 
-	return tp
+	return
 }
 
 func (p *TopicPublisher) Shutdown() error {

weed/mq/client/pub_client/scheduler.go

@@ -7,7 +7,9 @@ import (
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
 	"github.com/seaweedfs/seaweedfs/weed/util/buffered_queue"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/status"
 	"log"
 	"sort"
@@ -33,6 +35,7 @@ type EachPartitionPublishJob struct {
 func (p *TopicPublisher) startSchedulerThread(wg *sync.WaitGroup) error {
 
 	if err := p.doConfigureTopic(); err != nil {
+		wg.Done()
 		return fmt.Errorf("configure topic %s: %v", p.config.Topic, err)
 	}
 
@@ -111,6 +114,7 @@ func (p *TopicPublisher) onEachAssignments(generation int, assignments []*mq_pb.
 		go func(job *EachPartitionPublishJob) {
 			defer job.wg.Done()
 			if err := p.doPublishToPartition(job); err != nil {
+				log.Printf("publish to %s partition %v: %v", p.config.Topic, job.Partition, err)
 				errChan <- EachPartitionError{assignment, err, generation}
 			}
 		}(job)
@@ -126,7 +130,7 @@ func (p *TopicPublisher) doPublishToPartition(job *EachPartitionPublishJob) erro
 
 	log.Printf("connecting to %v for topic partition %+v", job.LeaderBroker, job.Partition)
 
-	grpcConnection, err := pb.GrpcDial(context.Background(), job.LeaderBroker, true, p.grpcDialOption)
+	grpcConnection, err := grpc.NewClient(job.LeaderBroker, grpc.WithTransportCredentials(insecure.NewCredentials()), p.grpcDialOption)
 	if err != nil {
 		return fmt.Errorf("dial broker %s: %v", job.LeaderBroker, err)
 	}
@@ -225,7 +229,7 @@ func (p *TopicPublisher) doPublishToPartition(job *EachPartitionPublishJob) erro
 
 func (p *TopicPublisher) doConfigureTopic() (err error) {
 	if len(p.config.Brokers) == 0 {
-		return fmt.Errorf("no bootstrap brokers")
+		return fmt.Errorf("topic configuring found no bootstrap brokers")
 	}
 	var lastErr error
 	for _, brokerAddress := range p.config.Brokers {
@@ -256,7 +260,7 @@ func (p *TopicPublisher) doConfigureTopic() (err error) {
 
 func (p *TopicPublisher) doLookupTopicPartitions() (assignments []*mq_pb.BrokerPartitionAssignment, err error) {
 	if len(p.config.Brokers) == 0 {
-		return nil, fmt.Errorf("no bootstrap brokers")
+		return nil, fmt.Errorf("lookup found no bootstrap brokers")
 	}
 	var lastErr error
 	for _, brokerAddress := range p.config.Brokers {

weed/mq/client/sub_client/connect_to_sub_coordinator.go

@@ -1,7 +1,6 @@
 package sub_client
 
 import (
-	"context"
 	"github.com/seaweedfs/seaweedfs/weed/glog"
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
@@ -12,10 +11,17 @@ func (sub *TopicSubscriber) doKeepConnectedToSubCoordinator() {
 	waitTime := 1 * time.Second
 	for {
 		for _, broker := range sub.bootstrapBrokers {
+
+			select {
+			case <-sub.ctx.Done():
+				return
+			default:
+			}
+
 			// lookup topic brokers
 			var brokerLeader string
 			err := pb.WithBrokerGrpcClient(false, broker, sub.SubscriberConfig.GrpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {
-				resp, err := client.FindBrokerLeader(context.Background(), &mq_pb.FindBrokerLeaderRequest{})
+				resp, err := client.FindBrokerLeader(sub.ctx, &mq_pb.FindBrokerLeaderRequest{})
 				if err != nil {
 					return err
 				}
@@ -30,10 +36,8 @@ func (sub *TopicSubscriber) doKeepConnectedToSubCoordinator() {
 
 			// connect to the balancer
 			pb.WithBrokerGrpcClient(true, brokerLeader, sub.SubscriberConfig.GrpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {
-				ctx, cancel := context.WithCancel(context.Background())
-				defer cancel()
 
-				stream, err := client.SubscriberToSubCoordinator(ctx)
+				stream, err := client.SubscriberToSubCoordinator(sub.ctx)
 				if err != nil {
 					glog.V(0).Infof("subscriber %s: %v", sub.ContentConfig.Topic, err)
 					return err
@@ -58,6 +62,13 @@ func (sub *TopicSubscriber) doKeepConnectedToSubCoordinator() {
 
 				go func() {
 					for reply := range sub.brokerPartitionAssignmentAckChan {
+
+						select {
+						case <-sub.ctx.Done():
+							return
+						default:
+						}
+
 						glog.V(0).Infof("subscriber instance %s ack %+v", sub.SubscriberConfig.ConsumerGroupInstanceId, reply)
 						if err := stream.Send(reply); err != nil {
 							glog.V(0).Infof("subscriber %s reply: %v", sub.ContentConfig.Topic, err)
@@ -73,6 +84,13 @@ func (sub *TopicSubscriber) doKeepConnectedToSubCoordinator() {
 						glog.V(0).Infof("subscriber %s receive: %v", sub.ContentConfig.Topic, err)
 						return err
 					}
+
+					select {
+					case <-sub.ctx.Done():
+						return nil
+					default:
+					}
+
 					sub.brokerPartitionAssignmentChan <- resp
 					glog.V(0).Infof("Received assignment: %+v", resp)
 				}

weed/mq/client/sub_client/on_each_partition.go

@@ -2,14 +2,13 @@ package sub_client
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"github.com/seaweedfs/seaweedfs/weed/glog"
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
 	"io"
-	"reflect"
-	"time"
 )
 
 type KeyedOffset struct {
@@ -35,8 +34,7 @@ func (sub *TopicSubscriber) onEachPartition(assigned *mq_pb.BrokerPartitionAssig
 		if po == nil {
 			po = &schema_pb.PartitionOffset{
 				Partition: assigned.Partition,
-				StartTsNs: time.Now().UnixNano(),
-				StartType: schema_pb.PartitionOffsetStartType_EARLIEST_IN_MEMORY,
+				StartTsNs: sub.ContentConfig.OffsetTsNs,
 			}
 		}
 
@@ -47,6 +45,7 @@ func (sub *TopicSubscriber) onEachPartition(assigned *mq_pb.BrokerPartitionAssig
 					ConsumerId:        sub.SubscriberConfig.ConsumerGroupInstanceId,
 					Topic:             sub.ContentConfig.Topic.ToPbTopic(),
 					PartitionOffset:   po,
+					OffsetType:        sub.ContentConfig.OffsetType,
 					Filter:            sub.ContentConfig.Filter,
 					FollowerBroker:    assigned.FollowerBroker,
 					SlidingWindowSize: slidingWindowSize,
@@ -65,6 +64,9 @@ func (sub *TopicSubscriber) onEachPartition(assigned *mq_pb.BrokerPartitionAssig
 		go func() {
 			for {
 				select {
+				case <-sub.ctx.Done():
+					subscribeClient.CloseSend()
+					return
 				case <-stopCh:
 					subscribeClient.CloseSend()
 					return
@@ -86,15 +88,27 @@ func (sub *TopicSubscriber) onEachPartition(assigned *mq_pb.BrokerPartitionAssig
 		}()
 
 		for {
-			// glog.V(0).Infof("subscriber %s/%s/%s waiting for message", sub.ContentConfig.Namespace, sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup)
+			// glog.V(0).Infof("subscriber %s/%s waiting for message", sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup)
 			resp, err := subscribeClient.Recv()
 			if err != nil {
+				if errors.Is(err, io.EOF) {
+					return nil
+				}
 				return fmt.Errorf("subscribe recv: %v", err)
 			}
 			if resp.Message == nil {
 				glog.V(0).Infof("subscriber %s/%s received nil message", sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup)
 				continue
 			}
+
+			select {
+			case <-sub.ctx.Done():
+				return nil
+			case <-stopCh:
+				return nil
+			default:
+			}
+
 			switch m := resp.Message.(type) {
 			case *mq_pb.SubscribeMessageResponse_Data:
 				if m.Data.Ctrl != nil {
@@ -102,7 +116,7 @@ func (sub *TopicSubscriber) onEachPartition(assigned *mq_pb.BrokerPartitionAssig
 					continue
 				}
 				if len(m.Data.Key) == 0 {
-					fmt.Printf("empty key %+v, type %v\n", m, reflect.TypeOf(m))
+					// fmt.Printf("empty key %+v, type %v\n", m, reflect.TypeOf(m))
 					continue
 				}
 				onDataMessageFn(m)

weed/mq/client/sub_client/subscribe.go

@@ -72,12 +72,12 @@ func (sub *TopicSubscriber) startProcessors() {
 				executors := util.NewLimitedConcurrentExecutor(int(sub.SubscriberConfig.SlidingWindowSize))
 				onDataMessageFn := func(m *mq_pb.SubscribeMessageResponse_Data) {
 					executors.Execute(func() {
-						processErr := sub.OnEachMessageFunc(m.Data.Key, m.Data.Value)
-						if processErr == nil {
-							sub.PartitionOffsetChan <- KeyedOffset{
-								Key:    m.Data.Key,
-								Offset: m.Data.TsNs,
-							}
+						if sub.OnDataMessageFunc != nil {
+							sub.OnDataMessageFunc(m)
+						}
+						sub.PartitionOffsetChan <- KeyedOffset{
+							Key:    m.Data.Key,
+							Offset: m.Data.TsNs,
 						}
 					})
 				}

weed/mq/client/sub_client/subscriber.go

@@ -1,6 +1,7 @@
 package sub_client
 
 import (
+	"context"
 	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
@@ -17,50 +18,50 @@ type SubscriberConfiguration struct {
 	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 OnEachMessageFunc func(key, value []byte) (err error)
 type OnCompletionFunc func()
 
 type TopicSubscriber struct {
+	ctx                              context.Context
 	SubscriberConfig                 *SubscriberConfiguration
 	ContentConfig                    *ContentConfiguration
 	brokerPartitionAssignmentChan    chan *mq_pb.SubscriberToSubCoordinatorResponse
 	brokerPartitionAssignmentAckChan chan *mq_pb.SubscriberToSubCoordinatorRequest
-	OnDataMessageFnnc                OnDataMessageFn
-	OnEachMessageFunc                OnEachMessageFunc
+	OnDataMessageFunc                OnDataMessageFn
 	OnCompletionFunc                 OnCompletionFunc
 	bootstrapBrokers                 []string
-	waitForMoreMessage               bool
 	activeProcessors                 map[topic.Partition]*ProcessorState
 	activeProcessorsLock             sync.Mutex
 	PartitionOffsetChan              chan KeyedOffset
 }
 
-func NewTopicSubscriber(bootstrapBrokers []string, subscriber *SubscriberConfiguration, content *ContentConfiguration, partitionOffsetChan chan KeyedOffset) *TopicSubscriber {
+func NewTopicSubscriber(ctx context.Context, bootstrapBrokers []string, subscriber *SubscriberConfiguration, content *ContentConfiguration, partitionOffsetChan chan KeyedOffset) *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,
-		waitForMoreMessage:               true,
 		activeProcessors:                 make(map[topic.Partition]*ProcessorState),
 		PartitionOffsetChan:              partitionOffsetChan,
 	}
 }
 
-func (sub *TopicSubscriber) SetEachMessageFunc(onEachMessageFn OnEachMessageFunc) {
-	sub.OnEachMessageFunc = onEachMessageFn
-}
-
 func (sub *TopicSubscriber) SetOnDataMessageFn(fn OnDataMessageFn) {
-	sub.OnDataMessageFnnc = fn
+	sub.OnDataMessageFunc = fn
 }
 
 func (sub *TopicSubscriber) SetCompletionFunc(onCompletionFn OnCompletionFunc) {

weed/mq/logstore/read_log_from_disk.go

@@ -36,7 +36,7 @@ func GenLogOnDiskReadFunc(filerClient filer_pb.FilerClient, t topic.Topic, p top
 				err = fmt.Errorf("unexpected unmarshal mq_pb.Message: %v", err)
 				return
 			}
-			if logEntry.TsNs < starTsNs {
+			if logEntry.TsNs <= starTsNs {
 				pos += 4 + int(size)
 				continue
 			}

weed/mq/logstore/read_parquet_to_log.go

@@ -73,7 +73,7 @@ func GenParquetReadFunc(filerClient filer_pb.FilerClient, t topic.Topic, p topic
 					return processedTsNs, fmt.Errorf("ToRecordValue failed: %v", err)
 				}
 				processedTsNs = recordValue.Fields[SW_COLUMN_NAME_TS].GetInt64Value()
-				if processedTsNs < starTsNs {
+				if processedTsNs <= starTsNs {
 					continue
 				}
 				if stopTsNs != 0 && processedTsNs >= stopTsNs {

weed/mq/sub_coordinator/market.go

@@ -20,7 +20,7 @@ Trigger rebalance when:
 * A new consumer group instance is added
 * Some partitions are unassigned from a consumer group instance.
 
-If multiple reblance requests are received, after a certain period, the market will only process the latest request.
+If multiple rebalance requests are received, after a certain period, the market will only process the latest request.
 
 However, if the number of unassigned partition is increased to exactly the total number of partitions,
 and total partitions are less than or equal to the sum of the max partition count of all consumer group instances,

weed/mq/topic/local_partition.go

@@ -89,12 +89,16 @@ func (p *LocalPartition) Subscribe(clientName string, startPosition log_buffer.M
 			return nil
 		}
 
-		startPosition = processedPosition
+		if processedPosition.Time.UnixNano() != 0 {
+			startPosition = processedPosition
+		}
 		processedPosition, isDone, readInMemoryLogErr = p.LogBuffer.LoopProcessLogData(clientName, startPosition, 0, onNoMessageFn, eachMessageFn)
 		if isDone {
 			return nil
 		}
-		startPosition = processedPosition
+		if processedPosition.Time.UnixNano() != 0 {
+			startPosition = processedPosition
+		}
 
 		if readInMemoryLogErr == log_buffer.ResumeFromDiskError {
 			continue