Squashed commit of the following:

commit 32f4b1a13057d56b6de487cdb80ff7c205af01a6
Author: chrislu <chris.lu@gmail.com>
Date:   Sun Aug 20 22:52:19 2023 -0700

    fix compilation

commit e77ad33b7ca0423138fbae26a4433b60923a9588
Author: chrislu <chris.lu@gmail.com>
Date:   Sun Aug 20 22:46:44 2023 -0700

    pub

commit f431f30cc7ca277ca299e3cd118c05537fb9f5c3
Author: chrislu <chris.lu@gmail.com>
Date:   Sun Aug 20 13:27:39 2023 -0700

    fix generic type

commit 4e9dcb18293fd1e3e306e2dceb995dfd67a35e1d
Merge: 30f942580 16e3f2d52
Author: chrislu <chris.lu@gmail.com>
Date:   Sun Aug 20 12:47:14 2023 -0700

    Merge branch 'master' into pubsub

commit 30f942580ad1bb32ae94aade2e3a21ec3ab63e21
Author: chrislu <chris.lu@gmail.com>
Date:   Sun Aug 20 11:10:58 2023 -0700

    wip

commit f8b00980bc2f3879bb43decffd9a08d842f196f2
Author: chrislu <chris.lu@gmail.com>
Date:   Tue Jul 25 09:14:35 2023 -0700

    add design document

commit 08d2bebe42a26ebc39f1542f54d99e73620727dd
Author: chrislu <chris.lu@gmail.com>
Date:   Tue Jul 25 09:14:06 2023 -0700

    minor

commit bcfa7982b262a40fcdce6fc6613fad2ce07c13da
Author: chrislu <chris.lu@gmail.com>
Date:   Tue Jul 25 09:13:49 2023 -0700

    rename

go.mod

@@ -249,6 +249,8 @@ require (
 	github.com/ncw/swift/v2 v2.0.1 // indirect
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/oracle/oci-go-sdk/v65 v65.34.0 // indirect
+	github.com/orcaman/concurrent-map v1.0.0 // indirect
+	github.com/orcaman/concurrent-map/v2 v2.0.1 // indirect
 	github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
 	github.com/pelletier/go-toml/v2 v2.0.8 // indirect
 	github.com/pengsrc/go-shared v0.2.1-0.20190131101655-1999055a4a14 // indirect

go.sum

@@ -638,6 +638,12 @@ github.com/opentracing/opentracing-go v1.2.0 h1:uEJPy/1a5RIPAJ0Ov+OIO8OxWu77jEv+
 github.com/opentracing/opentracing-go v1.2.0/go.mod h1:GxEUsuufX4nBwe+T+Wl9TAgYrxe9dPLANfrWvHYVTgc=
 github.com/oracle/oci-go-sdk/v65 v65.34.0 h1:uG1KucBxAbn8cYRgQHxtQKogtl85nOX8LhimZCPfMqw=
 github.com/oracle/oci-go-sdk/v65 v65.34.0/go.mod h1:MXMLMzHnnd9wlpgadPkdlkZ9YrwQmCOmbX5kjVEJodw=
+github.com/orcaman/concurrent-map v1.0.0 h1:I/2A2XPCb4IuQWcQhBhSwGfiuybl/J0ev9HDbW65HOY=
+github.com/orcaman/concurrent-map v1.0.0/go.mod h1:Lu3tH6HLW3feq74c2GC+jIMS/K2CFcDWnWD9XkenwhI=
+github.com/orcaman/concurrent-map/v2 v2.0.1 h1:jOJ5Pg2w1oeB6PeDurIYf6k9PQ+aTITr/6lP/L/zp6c=
+github.com/orcaman/concurrent-map/v2 v2.0.1/go.mod h1:9Eq3TG2oBe5FirmYWQfYO5iH1q0Jv47PLaNK++uCdOM=
+github.com/ovh/go-ovh v1.4.1/go.mod h1:6bL6pPyUT7tBfI0pqOegJgRjgjuO+mOo+MyXd1EEC0M=
+github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
 github.com/pascaldekloe/goe v0.1.0 h1:cBOtyMzM9HTpWjXfbbunk26uA6nG3a8n06Wieeh0MwY=
 github.com/pascaldekloe/goe v0.1.0/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
 github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=

weed/mq/broker/brokder_grpc_admin.go

@@ -4,7 +4,7 @@ import (
 	"context"
 	"github.com/seaweedfs/seaweedfs/weed/cluster"
 	"github.com/seaweedfs/seaweedfs/weed/glog"
-	"github.com/seaweedfs/seaweedfs/weed/mq"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
@@ -12,6 +12,10 @@ import (
 	"sync"
 )
 
+const (
+	MaxPartitionCount = 1024
+)
+
 func (broker *MessageQueueBroker) FindBrokerLeader(c context.Context, request *mq_pb.FindBrokerLeaderRequest) (*mq_pb.FindBrokerLeaderResponse, error) {
 	ret := &mq_pb.FindBrokerLeaderResponse{}
 	err := broker.withMasterClient(false, broker.MasterClient.GetMaster(), func(client master_pb.SeaweedClient) error {
@@ -31,21 +35,9 @@ func (broker *MessageQueueBroker) FindBrokerLeader(c context.Context, request *m
 	return ret, err
 }
 
-func (broker *MessageQueueBroker) CheckSegmentStatus(c context.Context, request *mq_pb.CheckSegmentStatusRequest) (*mq_pb.CheckSegmentStatusResponse, error) {
-	ret := &mq_pb.CheckSegmentStatusResponse{}
-	// TODO add in memory active segment
-	return ret, nil
-}
-
-func (broker *MessageQueueBroker) CheckBrokerLoad(c context.Context, request *mq_pb.CheckBrokerLoadRequest) (*mq_pb.CheckBrokerLoadResponse, error) {
-	ret := &mq_pb.CheckBrokerLoadResponse{}
-	// TODO read broker's load
-	return ret, nil
-}
-
 func (broker *MessageQueueBroker) AssignSegmentBrokers(c context.Context, request *mq_pb.AssignSegmentBrokersRequest) (*mq_pb.AssignSegmentBrokersResponse, error) {
 	ret := &mq_pb.AssignSegmentBrokersResponse{}
-	segment := mq.FromPbSegment(request.Segment)
+	segment := topic.FromPbSegment(request.Segment)
 
 	// check existing segment locations on filer
 	existingBrokers, err := broker.checkSegmentOnFiler(segment)
@@ -84,7 +76,92 @@ func (broker *MessageQueueBroker) AssignSegmentBrokers(c context.Context, reques
 	return ret, nil
 }
 
-func (broker *MessageQueueBroker) checkSegmentsOnBrokers(segment *mq.Segment, brokers []pb.ServerAddress) (active bool, err error) {
+func (broker *MessageQueueBroker) CheckSegmentStatus(c context.Context, request *mq_pb.CheckSegmentStatusRequest) (*mq_pb.CheckSegmentStatusResponse, error) {
+	ret := &mq_pb.CheckSegmentStatusResponse{}
+	// TODO add in memory active segment
+	return ret, nil
+}
+
+func (broker *MessageQueueBroker) CheckBrokerLoad(c context.Context, request *mq_pb.CheckBrokerLoadRequest) (*mq_pb.CheckBrokerLoadResponse, error) {
+	ret := &mq_pb.CheckBrokerLoadResponse{}
+	// TODO read broker's load
+	return ret, nil
+}
+
+// FindTopicBrokers returns the brokers that are serving the topic
+//
+//  1. lock the topic
+//
+//  2. find the topic partitions on the filer
+//     2.1 if the topic is not found, return error
+//     2.2 if the request is_for_publish, create the topic
+//     2.2.1 if the request is_for_subscribe, return error not found
+//     2.2.2 if the request is_for_publish, create the topic
+//     2.2 if the topic is found, return the brokers
+//
+//  3. unlock the topic
+func (broker *MessageQueueBroker) FindTopicBrokers(c context.Context, request *mq_pb.FindTopicBrokersRequest) (*mq_pb.FindTopicBrokersResponse, error) {
+	ret := &mq_pb.FindTopicBrokersResponse{}
+	// lock the topic
+
+	// find the topic partitions on the filer
+	// if the topic is not found
+	//   if the request is_for_publish
+	//     create the topic
+	//   if the request is_for_subscribe
+	//     return error not found
+	return ret, nil
+}
+
+// CheckTopicPartitionsStatus check the topic partitions on the broker
+func (broker *MessageQueueBroker) CheckTopicPartitionsStatus(c context.Context, request *mq_pb.CheckTopicPartitionsStatusRequest) (*mq_pb.CheckTopicPartitionsStatusResponse, error) {
+	ret := &mq_pb.CheckTopicPartitionsStatusResponse{}
+	return ret, nil
+}
+
+// createOrUpdateTopicPartitions creates the topic partitions on the broker
+// 1. check
+func (broker *MessageQueueBroker) createOrUpdateTopicPartitions(topic *topic.Topic, prevAssignment *mq_pb.TopicPartitionsAssignment) (err error) {
+	// create or update each partition
+	if prevAssignment == nil {
+		broker.createOrUpdateTopicPartition(topic, nil)
+	} else {
+		for _, partitionAssignment := range prevAssignment.BrokerPartitions {
+			broker.createOrUpdateTopicPartition(topic, partitionAssignment)
+		}
+	}
+	return nil
+}
+
+func (broker *MessageQueueBroker) createOrUpdateTopicPartition(topic *topic.Topic, oldAssignment *mq_pb.BrokerPartitionsAssignment) (newAssignment *mq_pb.BrokerPartitionsAssignment) {
+	shouldCreate := broker.confirmBrokerPartitionAssignment(topic, oldAssignment)
+	if !shouldCreate {
+
+	}
+	return
+}
+func (broker *MessageQueueBroker) confirmBrokerPartitionAssignment(topic *topic.Topic, oldAssignment *mq_pb.BrokerPartitionsAssignment) (shouldCreate bool) {
+	if oldAssignment == nil {
+		return true
+	}
+	for _, b := range oldAssignment.FollowerBrokers {
+		pb.WithBrokerClient(false, pb.ServerAddress(b), broker.grpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {
+			_, err := client.CheckTopicPartitionsStatus(context.Background(), &mq_pb.CheckTopicPartitionsStatusRequest{
+				Namespace:                  string(topic.Namespace),
+				Topic:                      topic.Name,
+				BrokerPartitionsAssignment: oldAssignment,
+				ShouldCancelIfNotMatch:     true,
+			})
+			if err != nil {
+				shouldCreate = true
+			}
+			return nil
+		})
+	}
+	return
+}
+
+func (broker *MessageQueueBroker) checkSegmentsOnBrokers(segment *topic.Segment, brokers []pb.ServerAddress) (active bool, err error) {
 	var wg sync.WaitGroup
 
 	for _, candidate := range brokers {

weed/mq/broker/brokder_grpc_pub.go

@@ -1,16 +1,136 @@
 package broker
 
 import (
+	"context"
+	"fmt"
+	"github.com/seaweedfs/seaweedfs/weed/glog"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
 )
 
+// For a new or re-configured topic, or one of the broker went offline,
+//   the pub clients ask one broker what are the brokers for all the topic partitions.
+// The broker will lock the topic on write.
+//   1. if the topic is not found, create the topic, and allocate the topic partitions to the brokers
+//   2. if the topic is found, return the brokers for the topic partitions
+// For a topic to read from, the sub clients ask one broker what are the brokers for all the topic partitions.
+// The broker will lock the topic on read.
+//   1. if the topic is not found, return error
+//   2. if the topic is found, return the brokers for the topic partitions
+//
+// If the topic needs to be re-balanced, the admin client will lock the topic,
+// 1. collect throughput information for all the brokers
+// 2. adjust the topic partitions to the brokers
+// 3. notify the brokers to add/remove partitions to host
+//    3.1 When locking the topic, the partitions and brokers should be remembered in the lock.
+// 4. the brokers will stop process incoming messages if not the right partition
+//    4.1 the pub clients will need to re-partition the messages and publish to the right brokers for the partition3
+//    4.2 the sub clients will need to change the brokers to read from
+//
+// The following is from each individual component's perspective:
+// For a pub client
+//   For current topic/partition, ask one broker for the brokers for the topic partitions
+//     1. connect to the brokers and keep sending, until the broker returns error, or the broker leader is moved.
+// For a sub client
+//   For current topic/partition, ask one broker for the brokers for the topic partitions
+//     1. connect to the brokers and keep reading, until the broker returns error, or the broker leader is moved.
+// For a broker
+//   Upon a pub client lookup:
+//     1. lock the topic
+//       2. if already has topic partition assignment, check all brokers are healthy
+//       3. if not, create topic partition assignment
+//     2. return the brokers for the topic partitions
+//     3. unlock the topic
+//   Upon a sub client lookup:
+//     1. lock the topic
+//       2. if already has topic partition assignment, check all brokers are healthy
+//       3. if not, return error
+//     2. return the brokers for the topic partitions
+//     3. unlock the topic
+// For an admin tool
+//   0. collect stats from all the brokers, and find the topic worth moving
+//   1. lock the topic
+//   2. collect throughput information for all the brokers
+//   3. adjust the topic partitions to the brokers
+//   4. notify the brokers to add/remove partitions to host
+//   5. the brokers will stop process incoming messages if not the right partition
+//   6. unlock the topic
+
 /*
-The messages is buffered in memory, and saved to filer under
+The messages are buffered in memory, and saved to filer under
 	/topics/<topic>/<date>/<hour>/<segment>/*.msg
 	/topics/<topic>/<date>/<hour>/segment
 	/topics/<topic>/info/segment_<id>.meta
+
+
+
 */
 
 func (broker *MessageQueueBroker) Publish(stream mq_pb.SeaweedMessaging_PublishServer) error {
+	// 1. write to the volume server
+	// 2. find the topic metadata owning filer
+	// 3. write to the filer
+
+	var localTopicPartition *topic.LocalPartition
+	for {
+		req, err := stream.Recv()
+		if err != nil {
+			return err
+		}
+
+		// Process the received message
+		sequence := req.GetSequence()
+		response := &mq_pb.PublishResponse{
+			AckSequence: sequence,
+		}
+		if dataMessage := req.GetData(); dataMessage != nil {
+			if localTopicPartition == nil {
+				response.Error = "topic partition not initialized"
+				glog.Errorf("topic partition not found")
+			} else {
+				localTopicPartition.Publish(dataMessage)
+			}
+		} else if initMessage := req.GetInit(); initMessage != nil {
+			localTopicPartition = broker.localTopicManager.GetTopicPartition(
+				topic.NewTopic(topic.Namespace(initMessage.Segment.Namespace), initMessage.Segment.Topic),
+				topic.FromPbPartition(initMessage.Segment.Partition),
+			)
+			if localTopicPartition == nil {
+				response.Error = fmt.Sprintf("topic partition %v not found", initMessage.Segment)
+				glog.Errorf("topic partition %v not found", initMessage.Segment)
+			}
+		}
+		if err := stream.Send(response); err != nil {
+			glog.Errorf("Error sending setup response: %v", err)
+		}
+	}
+
 	return nil
 }
+
+// AssignTopicPartitions Runs on the assigned broker, to execute the topic partition assignment
+func (broker *MessageQueueBroker) AssignTopicPartitions(c context.Context, request *mq_pb.AssignTopicPartitionsRequest) (*mq_pb.AssignTopicPartitionsResponse, error) {
+	ret := &mq_pb.AssignTopicPartitionsResponse{}
+	self := pb.ServerAddress(fmt.Sprintf("%s:%d", broker.option.Ip, broker.option.Port))
+
+	for _, partition := range request.TopicPartitionsAssignment.BrokerPartitions {
+		localPartiton := topic.FromPbBrokerPartitionsAssignment(self, partition)
+		broker.localTopicManager.AddTopicPartition(
+			topic.FromPbTopic(request.Topic),
+			localPartiton)
+		if request.IsLeader {
+			for _, follower := range localPartiton.FollowerBrokers {
+				err := pb.WithBrokerClient(false, follower, broker.grpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {
+					_, err := client.AssignTopicPartitions(context.Background(), request)
+					return err
+				})
+				if err != nil {
+					return ret, err
+				}
+			}
+		}
+	}
+
+	return ret, nil
+}

weed/mq/broker/broker_segment_serde.go

@@ -4,7 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"github.com/seaweedfs/seaweedfs/weed/filer"
-	"github.com/seaweedfs/seaweedfs/weed/mq"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/filer_pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
@@ -12,8 +12,8 @@ import (
 	"time"
 )
 
-func (broker *MessageQueueBroker) checkSegmentOnFiler(segment *mq.Segment) (brokers []pb.ServerAddress, err error) {
+func (broker *MessageQueueBroker) checkSegmentOnFiler(segment *topic.Segment) (brokers []pb.ServerAddress, err error) {
-	info, found, err := broker.readSegmentOnFiler(segment)
+	info, found, err := broker.readSegmentInfoOnFiler(segment)
 	if err != nil {
 		return
 	}
@@ -27,12 +27,12 @@ func (broker *MessageQueueBroker) checkSegmentOnFiler(segment *mq.Segment) (brok
 	return
 }
 
-func (broker *MessageQueueBroker) saveSegmentBrokersOnFiler(segment *mq.Segment, brokers []pb.ServerAddress) (err error) {
+func (broker *MessageQueueBroker) saveSegmentBrokersOnFiler(segment *topic.Segment, brokers []pb.ServerAddress) (err error) {
 	var nodes []string
 	for _, b := range brokers {
 		nodes = append(nodes, string(b))
 	}
-	broker.saveSegmentToFiler(segment, &mq_pb.SegmentInfo{
+	broker.saveSegmentInfoToFiler(segment, &mq_pb.SegmentInfo{
 		Segment:          segment.ToPbSegment(),
 		StartTsNs:        time.Now().UnixNano(),
 		Brokers:          nodes,
@@ -43,7 +43,7 @@ func (broker *MessageQueueBroker) saveSegmentBrokersOnFiler(segment *mq.Segment,
 	return
 }
 
-func (broker *MessageQueueBroker) readSegmentOnFiler(segment *mq.Segment) (info *mq_pb.SegmentInfo, found bool, err error) {
+func (broker *MessageQueueBroker) readSegmentInfoOnFiler(segment *topic.Segment) (info *mq_pb.SegmentInfo, found bool, err error) {
 	dir, name := segment.DirAndName()
 
 	found, err = filer_pb.Exists(broker, dir, name, false)
@@ -70,7 +70,7 @@ func (broker *MessageQueueBroker) readSegmentOnFiler(segment *mq.Segment) (info
 	return
 }
 
-func (broker *MessageQueueBroker) saveSegmentToFiler(segment *mq.Segment, info *mq_pb.SegmentInfo) (err error) {
+func (broker *MessageQueueBroker) saveSegmentInfoToFiler(segment *topic.Segment, info *mq_pb.SegmentInfo) (err error) {
 	dir, name := segment.DirAndName()
 
 	var buf bytes.Buffer

weed/mq/broker/broker_server.go

@@ -1,6 +1,7 @@
 package broker
 
 import (
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
 	"time"
 
 	"github.com/seaweedfs/seaweedfs/weed/cluster"
@@ -32,6 +33,7 @@ type MessageQueueBroker struct {
 	MasterClient      *wdclient.MasterClient
 	filers            map[pb.ServerAddress]struct{}
 	currentFiler      pb.ServerAddress
+	localTopicManager *topic.LocalTopicManager
 }
 
 func NewMessageBroker(option *MessageQueueBrokerOption, grpcDialOption grpc.DialOption) (mqBroker *MessageQueueBroker, err error) {
@@ -41,6 +43,7 @@ func NewMessageBroker(option *MessageQueueBrokerOption, grpcDialOption grpc.Dial
 		grpcDialOption:    grpcDialOption,
 		MasterClient:      wdclient.NewMasterClient(grpcDialOption, option.FilerGroup, cluster.BrokerType, pb.NewServerAddress(option.Ip, option.Port, 0), option.DataCenter, option.Rack, option.Masters),
 		filers:            make(map[pb.ServerAddress]struct{}),
+		localTopicManager: topic.NewLocalTopicManager(),
 	}
 	mqBroker.MasterClient.SetOnPeerUpdateFn(mqBroker.OnBrokerUpdate)
 

weed/mq/segment/message_serde_test.go

@@ -17,6 +17,7 @@ func TestMessageSerde(t *testing.T) {
 	bb := NewMessageBatchBuilder(b, 1, 2, 3, 4)
 
 	bb.AddMessage(5, 6, prop, []byte("the primary key"), []byte("body is here"))
+	bb.AddMessage(5, 7, prop, []byte("the primary 2"), []byte("body is 2"))
 
 	bb.BuildMessageBatch()
 
@@ -33,9 +34,9 @@ func TestMessageSerde(t *testing.T) {
 	assert.Equal(t, int64(5), mb.SegmentSeqBase())
 	assert.Equal(t, int32(0), mb.SegmentSeqMaxDelta())
 	assert.Equal(t, int64(6), mb.TsMsBase())
-	assert.Equal(t, int32(0), mb.TsMsMaxDelta())
+	assert.Equal(t, int32(1), mb.TsMsMaxDelta())
 
-	assert.Equal(t, 1, mb.MessagesLength())
+	assert.Equal(t, 2, mb.MessagesLength())
 
 	m := &message_fbs.Message{}
 	mb.Messages(m, 0)

weed/mq/topic/local_manager.go

@@ -0,0 +1,54 @@
+package topic
+
+import (
+	cmap "github.com/orcaman/concurrent-map/v2"
+)
+
+// LocalTopicManager manages topics on local broker
+type LocalTopicManager struct {
+	topics cmap.ConcurrentMap[string, *LocalTopic]
+}
+
+// NewLocalTopicManager creates a new LocalTopicManager
+func NewLocalTopicManager() *LocalTopicManager {
+	return &LocalTopicManager{
+		topics: cmap.New[*LocalTopic](),
+	}
+}
+
+// AddTopic adds a topic to the local topic manager
+func (manager *LocalTopicManager) AddTopicPartition(topic Topic, localPartition *LocalPartition) {
+	localTopic, ok := manager.topics.Get(topic.String())
+	if !ok {
+		localTopic = &LocalTopic{
+			Topic:      topic,
+			Partitions: make([]*LocalPartition, 0),
+		}
+	}
+	if localTopic.findPartition(localPartition.Partition) != nil {
+		return
+	}
+	localTopic.Partitions = append(localTopic.Partitions, localPartition)
+}
+
+// GetTopic gets a topic from the local topic manager
+func (manager *LocalTopicManager) GetTopicPartition(topic Topic, partition Partition) *LocalPartition {
+	localTopic, ok := manager.topics.Get(topic.String())
+	if !ok {
+		return nil
+	}
+	return localTopic.findPartition(partition)
+}
+
+// RemoveTopic removes a topic from the local topic manager
+func (manager *LocalTopicManager) RemoveTopic(topic Topic) {
+	manager.topics.Remove(topic.String())
+}
+
+func (manager *LocalTopicManager) RemoveTopicPartition(topic Topic, partition Partition) (removed bool) {
+	localTopic, ok := manager.topics.Get(topic.String())
+	if !ok {
+		return false
+	}
+	return localTopic.removePartition(partition)
+}

weed/mq/topic/local_partition.go

@@ -0,0 +1,40 @@
+package topic
+
+import (
+	"github.com/seaweedfs/seaweedfs/weed/pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/util/log_buffer"
+	"time"
+)
+
+type LocalPartition struct {
+	Partition
+	isLeader        bool
+	FollowerBrokers []pb.ServerAddress
+	logBuffer       *log_buffer.LogBuffer
+}
+
+func (p LocalPartition) Publish(message *mq_pb.PublishRequest_DataMessage) {
+	p.logBuffer.AddToBuffer(message.Key, message.Value, time.Now().UnixNano())
+}
+
+func FromPbBrokerPartitionsAssignment(self pb.ServerAddress, assignment *mq_pb.BrokerPartitionsAssignment) *LocalPartition {
+	isLeaer := assignment.LeaderBroker == string(self)
+	localPartition := &LocalPartition{
+		Partition: Partition{
+			RangeStart: assignment.PartitionStart,
+			RangeStop:  assignment.PartitionStop,
+			RingSize:   PartitionCount,
+		},
+		isLeader: isLeaer,
+	}
+	if !isLeaer {
+		return localPartition
+	}
+	followers := make([]pb.ServerAddress, len(assignment.FollowerBrokers))
+	for i, follower := range assignment.FollowerBrokers {
+		followers[i] = pb.ServerAddress(follower)
+	}
+	localPartition.FollowerBrokers = followers
+	return localPartition
+}

weed/mq/topic/local_topic.go

@@ -0,0 +1,29 @@
+package topic
+
+type LocalTopic struct {
+	Topic
+	Partitions []*LocalPartition
+}
+
+func (localTopic *LocalTopic) findPartition(partition Partition) *LocalPartition {
+	for _, localPartition := range localTopic.Partitions {
+		if localPartition.Partition.Equals(partition) {
+			return localPartition
+		}
+	}
+	return nil
+}
+func (localTopic *LocalTopic) removePartition(partition Partition) bool {
+	foundPartitionIndex := -1
+	for i, localPartition := range localTopic.Partitions {
+		if localPartition.Partition.Equals(partition) {
+			foundPartitionIndex = i
+			break
+		}
+	}
+	if foundPartitionIndex == -1 {
+		return false
+	}
+	localTopic.Partitions = append(localTopic.Partitions[:foundPartitionIndex], localTopic.Partitions[foundPartitionIndex+1:]...)
+	return true
+}

weed/mq/topic/partition.go

@@ -0,0 +1,32 @@
+package topic
+
+import "github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+
+const PartitionCount = 4096
+
+type Partition struct {
+	RangeStart int32
+	RangeStop  int32 // exclusive
+	RingSize   int32
+}
+
+func (partition Partition) Equals(other Partition) bool {
+	if partition.RangeStart != other.RangeStart {
+		return false
+	}
+	if partition.RangeStop != other.RangeStop {
+		return false
+	}
+	if partition.RingSize != other.RingSize {
+		return false
+	}
+	return true
+}
+
+func FromPbPartition(partition *mq_pb.Partition) Partition {
+	return Partition{
+		RangeStart: partition.RangeStart,
+		RangeStop:  partition.RangeStop,
+		RingSize:   partition.RingSize,
+	}
+}

weed/mq/topic/topic.go

@@ -1,4 +1,4 @@
-package mq
+package topic
 
 import (
 	"fmt"
@@ -14,10 +14,21 @@ type Topic struct {
 	Name      string
 }
 
-type Partition struct {
+func NewTopic(namespace Namespace, name string) Topic {
-	RangeStart int32
+	return Topic{
-	RangeStop  int32 // exclusive
+		Namespace: namespace,
-	RingSize   int32
+		Name:      name,
+	}
+}
+func FromPbTopic(topic *mq_pb.Topic) Topic {
+	return Topic{
+		Namespace: Namespace(topic.Namespace),
+		Name:      topic.Name,
+	}
+}
+
+func (tp Topic) String() string {
+	return fmt.Sprintf("%s.%s", tp.Namespace, tp.Name)
 }
 
 type Segment struct {

weed/mq/topic_allocation/allocation.go

@@ -0,0 +1,81 @@
+package topic_allocation
+
+import (
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/seaweedfs/seaweedfs/weed/pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"modernc.org/mathutil"
+)
+
+const (
+	DefaultBrokerCount = 4
+)
+
+// AllocateBrokersForTopicPartitions allocate brokers for a topic's all partitions
+func AllocateBrokersForTopicPartitions(t topic.Topic, prevAssignment *mq_pb.TopicPartitionsAssignment, candidateBrokers []pb.ServerAddress) (assignment *mq_pb.TopicPartitionsAssignment, err error) {
+	// create a previous assignment if not exists
+	if prevAssignment == nil || len(prevAssignment.BrokerPartitions) == 0 {
+		prevAssignment = &mq_pb.TopicPartitionsAssignment{
+			PartitionCount: topic.PartitionCount,
+		}
+		partitionCountForEachBroker := topic.PartitionCount / DefaultBrokerCount
+		for i := 0; i < DefaultBrokerCount; i++ {
+			prevAssignment.BrokerPartitions = append(prevAssignment.BrokerPartitions, &mq_pb.BrokerPartitionsAssignment{
+				PartitionStart: int32(i * partitionCountForEachBroker),
+				PartitionStop:  mathutil.MaxInt32(int32((i+1)*partitionCountForEachBroker), topic.PartitionCount),
+			})
+		}
+	}
+
+	// create a new assignment
+	assignment = &mq_pb.TopicPartitionsAssignment{
+		PartitionCount: prevAssignment.PartitionCount,
+	}
+
+	// allocate partitions for each partition range
+	for _, brokerPartition := range prevAssignment.BrokerPartitions {
+		// allocate partitions for each partition range
+		leader, followers, err := allocateBrokersForOneTopicPartition(t, brokerPartition, candidateBrokers)
+		if err != nil {
+			return nil, err
+		}
+
+		followerBrokers := make([]string, len(followers))
+		for i, follower := range followers {
+			followerBrokers[i] = string(follower)
+		}
+
+		assignment.BrokerPartitions = append(assignment.BrokerPartitions, &mq_pb.BrokerPartitionsAssignment{
+			PartitionStart:  brokerPartition.PartitionStart,
+			PartitionStop:   brokerPartition.PartitionStop,
+			LeaderBroker:    string(leader),
+			FollowerBrokers: followerBrokers,
+		})
+	}
+
+	return
+}
+
+func allocateBrokersForOneTopicPartition(t topic.Topic, brokerPartition *mq_pb.BrokerPartitionsAssignment, candidateBrokers []pb.ServerAddress) (leader pb.ServerAddress, followers []pb.ServerAddress, err error) {
+	// allocate leader
+	leader, err = allocateLeaderForOneTopicPartition(t, brokerPartition, candidateBrokers)
+	if err != nil {
+		return
+	}
+
+	// allocate followers
+	followers, err = allocateFollowersForOneTopicPartition(t, brokerPartition, candidateBrokers)
+	if err != nil {
+		return
+	}
+
+	return
+}
+
+func allocateFollowersForOneTopicPartition(t topic.Topic, partition *mq_pb.BrokerPartitionsAssignment, brokers []pb.ServerAddress) (followers []pb.ServerAddress, err error) {
+	return
+}
+
+func allocateLeaderForOneTopicPartition(t topic.Topic, partition *mq_pb.BrokerPartitionsAssignment, brokers []pb.ServerAddress) (leader pb.ServerAddress, err error) {
+	return
+}

weed/pb/mq.proto

@@ -20,6 +20,17 @@ service SeaweedMessaging {
     rpc CheckBrokerLoad (CheckBrokerLoadRequest) returns (CheckBrokerLoadResponse) {
     }
 
+    // control plane for topic partitions
+    rpc FindTopicBrokers (FindTopicBrokersRequest) returns (FindTopicBrokersResponse) {
+    }
+    // a pub client will call this to get the topic partitions assignment
+    rpc RequestTopicPartitions (RequestTopicPartitionsRequest) returns (RequestTopicPartitionsResponse) {
+    }
+    rpc AssignTopicPartitions (AssignTopicPartitionsRequest) returns (AssignTopicPartitionsResponse) {
+    }
+    rpc CheckTopicPartitionsStatus (CheckTopicPartitionsStatusRequest) returns (CheckTopicPartitionsStatusResponse) {
+    }
+
     // data plane
     rpc Publish (stream PublishRequest) returns (stream PublishResponse) {
     }
@@ -45,6 +56,10 @@ message FindBrokerLeaderResponse {
     string broker = 1;
 }
 
+message Topic {
+    string namespace = 1;
+    string name = 2;
+}
 message Partition {
     int32 ring_size = 1;
     int32 range_start = 2;
@@ -83,15 +98,68 @@ message CheckBrokerLoadResponse {
 
 }
 
+message FindTopicBrokersRequest {
+    Topic topic = 1;
+    bool is_for_publish = 2;
+}
+message FindTopicBrokersResponse {
+    Topic topic = 1;
+    TopicPartitionsAssignment topic_partitions_assignment = 2;
+}
+message BrokerPartitionsAssignment {
+    int32 partition_start = 1;
+    int32 partition_stop = 2;
+    string leader_broker = 3;
+    repeated string follower_brokers = 4;
+}
+message TopicPartitionsAssignment {
+    int32 partition_count = 1; // over-sharded partitions, usually 1024
+    repeated BrokerPartitionsAssignment broker_partitions = 2;
+}
+
+message RequestTopicPartitionsRequest {
+    Topic topic = 1;
+    int32 partition_count = 2;
+}
+message RequestTopicPartitionsResponse {
+    TopicPartitionsAssignment topic_partitions_assignment = 1;
+}
+
+message AssignTopicPartitionsRequest {
+    Topic topic = 1;
+    TopicPartitionsAssignment topic_partitions_assignment = 2;
+    bool is_leader = 3;
+}
+message AssignTopicPartitionsResponse {
+}
+
+message CheckTopicPartitionsStatusRequest {
+    string namespace = 1;
+    string topic = 2;
+    BrokerPartitionsAssignment broker_partitions_assignment = 3;
+    bool should_cancel_if_not_match = 4;
+}
+message CheckTopicPartitionsStatusResponse {
+    TopicPartitionsAssignment topic_partitions_assignment = 1;
+}
+
 //////////////////////////////////////////////////
 message PublishRequest {
     message InitMessage {
         Segment segment = 1;
     }
+    message DataMessage {
+        bytes key = 1;
+        bytes value = 2;
+    }
+    oneof message {
         InitMessage init = 1;
-    bytes message = 2;
+        DataMessage data = 2;
+    }
+    int64 sequence = 3;
 }
 message PublishResponse {
     int64 ack_sequence = 1;
-    bool is_closed = 2;
+    string error = 2;
+    bool is_closed = 3;
 }

weed/pb/mq_pb/mq_grpc.pb.go

@@ -27,6 +27,12 @@ type SeaweedMessagingClient interface {
 	AssignSegmentBrokers(ctx context.Context, in *AssignSegmentBrokersRequest, opts ...grpc.CallOption) (*AssignSegmentBrokersResponse, error)
 	CheckSegmentStatus(ctx context.Context, in *CheckSegmentStatusRequest, opts ...grpc.CallOption) (*CheckSegmentStatusResponse, error)
 	CheckBrokerLoad(ctx context.Context, in *CheckBrokerLoadRequest, opts ...grpc.CallOption) (*CheckBrokerLoadResponse, error)
+	// control plane for topic partitions
+	FindTopicBrokers(ctx context.Context, in *FindTopicBrokersRequest, opts ...grpc.CallOption) (*FindTopicBrokersResponse, error)
+	// a pub client will call this to get the topic partitions assignment
+	RequestTopicPartitions(ctx context.Context, in *RequestTopicPartitionsRequest, opts ...grpc.CallOption) (*RequestTopicPartitionsResponse, error)
+	AssignTopicPartitions(ctx context.Context, in *AssignTopicPartitionsRequest, opts ...grpc.CallOption) (*AssignTopicPartitionsResponse, error)
+	CheckTopicPartitionsStatus(ctx context.Context, in *CheckTopicPartitionsStatusRequest, opts ...grpc.CallOption) (*CheckTopicPartitionsStatusResponse, error)
 	// data plane
 	Publish(ctx context.Context, opts ...grpc.CallOption) (SeaweedMessaging_PublishClient, error)
 }
@@ -75,6 +81,42 @@ func (c *seaweedMessagingClient) CheckBrokerLoad(ctx context.Context, in *CheckB
 	return out, nil
 }
 
+func (c *seaweedMessagingClient) FindTopicBrokers(ctx context.Context, in *FindTopicBrokersRequest, opts ...grpc.CallOption) (*FindTopicBrokersResponse, error) {
+	out := new(FindTopicBrokersResponse)
+	err := c.cc.Invoke(ctx, "/messaging_pb.SeaweedMessaging/FindTopicBrokers", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *seaweedMessagingClient) RequestTopicPartitions(ctx context.Context, in *RequestTopicPartitionsRequest, opts ...grpc.CallOption) (*RequestTopicPartitionsResponse, error) {
+	out := new(RequestTopicPartitionsResponse)
+	err := c.cc.Invoke(ctx, "/messaging_pb.SeaweedMessaging/RequestTopicPartitions", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *seaweedMessagingClient) AssignTopicPartitions(ctx context.Context, in *AssignTopicPartitionsRequest, opts ...grpc.CallOption) (*AssignTopicPartitionsResponse, error) {
+	out := new(AssignTopicPartitionsResponse)
+	err := c.cc.Invoke(ctx, "/messaging_pb.SeaweedMessaging/AssignTopicPartitions", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *seaweedMessagingClient) CheckTopicPartitionsStatus(ctx context.Context, in *CheckTopicPartitionsStatusRequest, opts ...grpc.CallOption) (*CheckTopicPartitionsStatusResponse, error) {
+	out := new(CheckTopicPartitionsStatusResponse)
+	err := c.cc.Invoke(ctx, "/messaging_pb.SeaweedMessaging/CheckTopicPartitionsStatus", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
 func (c *seaweedMessagingClient) Publish(ctx context.Context, opts ...grpc.CallOption) (SeaweedMessaging_PublishClient, error) {
 	stream, err := c.cc.NewStream(ctx, &SeaweedMessaging_ServiceDesc.Streams[0], "/messaging_pb.SeaweedMessaging/Publish", opts...)
 	if err != nil {
@@ -115,6 +157,12 @@ type SeaweedMessagingServer interface {
 	AssignSegmentBrokers(context.Context, *AssignSegmentBrokersRequest) (*AssignSegmentBrokersResponse, error)
 	CheckSegmentStatus(context.Context, *CheckSegmentStatusRequest) (*CheckSegmentStatusResponse, error)
 	CheckBrokerLoad(context.Context, *CheckBrokerLoadRequest) (*CheckBrokerLoadResponse, error)
+	// control plane for topic partitions
+	FindTopicBrokers(context.Context, *FindTopicBrokersRequest) (*FindTopicBrokersResponse, error)
+	// a pub client will call this to get the topic partitions assignment
+	RequestTopicPartitions(context.Context, *RequestTopicPartitionsRequest) (*RequestTopicPartitionsResponse, error)
+	AssignTopicPartitions(context.Context, *AssignTopicPartitionsRequest) (*AssignTopicPartitionsResponse, error)
+	CheckTopicPartitionsStatus(context.Context, *CheckTopicPartitionsStatusRequest) (*CheckTopicPartitionsStatusResponse, error)
 	// data plane
 	Publish(SeaweedMessaging_PublishServer) error
 	mustEmbedUnimplementedSeaweedMessagingServer()
@@ -136,6 +184,18 @@ func (UnimplementedSeaweedMessagingServer) CheckSegmentStatus(context.Context, *
 func (UnimplementedSeaweedMessagingServer) CheckBrokerLoad(context.Context, *CheckBrokerLoadRequest) (*CheckBrokerLoadResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method CheckBrokerLoad not implemented")
 }
+func (UnimplementedSeaweedMessagingServer) FindTopicBrokers(context.Context, *FindTopicBrokersRequest) (*FindTopicBrokersResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method FindTopicBrokers not implemented")
+}
+func (UnimplementedSeaweedMessagingServer) RequestTopicPartitions(context.Context, *RequestTopicPartitionsRequest) (*RequestTopicPartitionsResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method RequestTopicPartitions not implemented")
+}
+func (UnimplementedSeaweedMessagingServer) AssignTopicPartitions(context.Context, *AssignTopicPartitionsRequest) (*AssignTopicPartitionsResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method AssignTopicPartitions not implemented")
+}
+func (UnimplementedSeaweedMessagingServer) CheckTopicPartitionsStatus(context.Context, *CheckTopicPartitionsStatusRequest) (*CheckTopicPartitionsStatusResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method CheckTopicPartitionsStatus not implemented")
+}
 func (UnimplementedSeaweedMessagingServer) Publish(SeaweedMessaging_PublishServer) error {
 	return status.Errorf(codes.Unimplemented, "method Publish not implemented")
 }
@@ -224,6 +284,78 @@ func _SeaweedMessaging_CheckBrokerLoad_Handler(srv interface{}, ctx context.Cont
 	return interceptor(ctx, in, info, handler)
 }
 
+func _SeaweedMessaging_FindTopicBrokers_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(FindTopicBrokersRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(SeaweedMessagingServer).FindTopicBrokers(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/messaging_pb.SeaweedMessaging/FindTopicBrokers",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(SeaweedMessagingServer).FindTopicBrokers(ctx, req.(*FindTopicBrokersRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+func _SeaweedMessaging_RequestTopicPartitions_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(RequestTopicPartitionsRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(SeaweedMessagingServer).RequestTopicPartitions(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/messaging_pb.SeaweedMessaging/RequestTopicPartitions",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(SeaweedMessagingServer).RequestTopicPartitions(ctx, req.(*RequestTopicPartitionsRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+func _SeaweedMessaging_AssignTopicPartitions_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(AssignTopicPartitionsRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(SeaweedMessagingServer).AssignTopicPartitions(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/messaging_pb.SeaweedMessaging/AssignTopicPartitions",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(SeaweedMessagingServer).AssignTopicPartitions(ctx, req.(*AssignTopicPartitionsRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+func _SeaweedMessaging_CheckTopicPartitionsStatus_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(CheckTopicPartitionsStatusRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(SeaweedMessagingServer).CheckTopicPartitionsStatus(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/messaging_pb.SeaweedMessaging/CheckTopicPartitionsStatus",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(SeaweedMessagingServer).CheckTopicPartitionsStatus(ctx, req.(*CheckTopicPartitionsStatusRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
 func _SeaweedMessaging_Publish_Handler(srv interface{}, stream grpc.ServerStream) error {
 	return srv.(SeaweedMessagingServer).Publish(&seaweedMessagingPublishServer{stream})
 }
@@ -273,6 +405,22 @@ var SeaweedMessaging_ServiceDesc = grpc.ServiceDesc{
 			MethodName: "CheckBrokerLoad",
 			Handler:    _SeaweedMessaging_CheckBrokerLoad_Handler,
 		},
+		{
+			MethodName: "FindTopicBrokers",
+			Handler:    _SeaweedMessaging_FindTopicBrokers_Handler,
+		},
+		{
+			MethodName: "RequestTopicPartitions",
+			Handler:    _SeaweedMessaging_RequestTopicPartitions_Handler,
+		},
+		{
+			MethodName: "AssignTopicPartitions",
+			Handler:    _SeaweedMessaging_AssignTopicPartitions_Handler,
+		},
+		{
+			MethodName: "CheckTopicPartitionsStatus",
+			Handler:    _SeaweedMessaging_CheckTopicPartitionsStatus_Handler,
+		},
 	},
 	Streams: []grpc.StreamDesc{
 		{