snapshot of https://github.com/seaweedfs/seaweedfs.git (master, 2026-09-16; blobs > 131072B omitted)

weed/mq/client/agent_client/publish_session.go

@@ -0,0 +1,82 @@
+package agent_client
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/seaweedfs/seaweedfs/weed/mq/schema"
+	"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 {
+	schema         *schema.Schema
+	partitionCount int
+	publisherName  string
+	stream         grpc.BidiStreamingClient[mq_agent_pb.PublishRecordRequest, mq_agent_pb.PublishRecordResponse]
+}
+
+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 := 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.StartPublishSession(context.Background(), &mq_agent_pb.StartPublishSessionRequest{
+		Topic: &schema_pb.Topic{
+			Namespace: topicSchema.Namespace,
+			Name:      topicSchema.Name,
+		},
+		PartitionCount: int32(partitionCount),
+		RecordType:     topicSchema.RecordType,
+		PublisherName:  publisherName,
+	})
+	if err != nil {
+		return nil, err
+	}
+	if resp.Error != "" {
+		return nil, fmt.Errorf("start publish session: %v", resp.Error)
+	}
+
+	stream, err := agentClient.PublishRecord(context.Background())
+	if err != nil {
+		return nil, fmt.Errorf("publish record: %w", err)
+	}
+
+	if err = stream.Send(&mq_agent_pb.PublishRecordRequest{
+		SessionId: resp.SessionId,
+	}); err != nil {
+		return nil, fmt.Errorf("send session id: %w", err)
+	}
+
+	return &PublishSession{
+		schema:         topicSchema,
+		partitionCount: partitionCount,
+		publisherName:  publisherName,
+		stream:         stream,
+	}, nil
+}
+
+func (a *PublishSession) CloseSession() error {
+	if a.schema == nil {
+		return nil
+	}
+	err := a.stream.CloseSend()
+	if err != nil {
+		return fmt.Errorf("close send: %w", err)
+	}
+	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

@@ -0,0 +1,87 @@
+package agent_client
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"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
+	SlidingWindowSize       int32
+}
+
+type SubscribeSession struct {
+	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 := 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)
+
+	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,
+		SlidingWindowSize:       option.SlidingWindowSize,
+	}
+
+	stream, err := agentClient.SubscribeRecord(context.Background())
+	if err != nil {
+		return nil, fmt.Errorf("subscribe record: %w", err)
+	}
+
+	if err = stream.Send(&mq_agent_pb.SubscribeRecordRequest{
+		Init: initRequest,
+	}); err != nil {
+		return nil, fmt.Errorf("send session id: %w", err)
+	}
+
+	return &SubscribeSession{
+		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 {
+			if onCompletionFn != nil {
+				onCompletionFn()
+			}
+			return err
+		}
+		onEachMessageFn(resp.Key, resp.Value)
+	}
+}