allocate brokers to serve segments

weed/mq/broker/brokder_grpc_admin.go

@@ -2,9 +2,211 @@ package broker
 
 import (
 	"context"
+	"github.com/chrislusf/seaweedfs/weed/cluster"
+	"github.com/chrislusf/seaweedfs/weed/glog"
+	"github.com/chrislusf/seaweedfs/weed/mq"
+	"github.com/chrislusf/seaweedfs/weed/pb"
+	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
 	"github.com/chrislusf/seaweedfs/weed/pb/mq_pb"
+	"sort"
+	"sync"
 )
 
 func (broker *MessageQueueBroker) FindBrokerLeader(c context.Context, request *mq_pb.FindBrokerLeaderRequest) (*mq_pb.FindBrokerLeaderResponse, error) {
-	panic("implement me")
+	ret := &mq_pb.FindBrokerLeaderResponse{}
+	err := broker.withMasterClient(false, broker.MasterClient.GetMaster(), func(client master_pb.SeaweedClient) error {
+		resp, err := client.ListClusterNodes(context.Background(), &master_pb.ListClusterNodesRequest{
+			ClientType:   cluster.BrokerType,
+			FilerGroup:   request.FilerGroup,
+			IsLeaderOnly: true,
+		})
+		if err != nil {
+			return err
+		}
+		if len(resp.ClusterNodes) == 0 {
+			return nil
+		}
+		ret.Broker = resp.ClusterNodes[0].Address
+		return nil
+	})
+	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)
+
+	// check existing segment locations on filer
+	existingBrokers, err := broker.checkSegmentOnFiler(segment)
+	if err != nil {
+		return ret, err
+	}
+	// good if the segment is still on the brokers
+	isActive, err := broker.checkSegmentsOnBrokers(segment, existingBrokers)
+	if err != nil {
+		return ret, err
+	}
+	if isActive {
+		for _, broker := range existingBrokers {
+			ret.Brokers = append(ret.Brokers, string(broker))
+		}
+		return ret, nil
+	}
+
+	// randomly pick up to 10 brokers, and find the ones with the lightest load
+	selectedBrokers, err := broker.selectBrokers()
+	if err != nil {
+		return ret, err
+	}
+
+	// save the allocated brokers info for this segment on the filer
+	if err := broker.saveSegmentOnFiler(segment, selectedBrokers); err != nil {
+		return ret, err
+	}
+
+	for _, broker := range selectedBrokers {
+		ret.Brokers = append(ret.Brokers, string(broker))
+	}
+	return ret, nil
+}
+
+func (broker *MessageQueueBroker) checkSegmentOnFiler(segment *mq.Segment) (brokers []pb.ServerAddress, err error) {
+	return
+}
+
+func (broker *MessageQueueBroker) checkSegmentsOnBrokers(segment *mq.Segment, brokers []pb.ServerAddress) (active bool, err error) {
+	var wg sync.WaitGroup
+
+	for _, candidate := range brokers {
+		wg.Add(1)
+		go func(candidate pb.ServerAddress) {
+			defer wg.Done()
+			broker.withBrokerClient(false, candidate, func(client mq_pb.SeaweedMessagingClient) error {
+				resp, checkErr := client.CheckSegmentStatus(context.Background(), &mq_pb.CheckSegmentStatusRequest{
+					Segment: &mq_pb.Segment{
+						Namespace: string(segment.Topic.Namespace),
+						Topic:     segment.Topic.Name,
+						Id:        segment.Id,
+					},
+				})
+				if checkErr != nil {
+					err = checkErr
+					glog.V(0).Infof("check segment status on %s: %v", candidate, checkErr)
+					return nil
+				}
+				if resp.IsActive == false {
+					active = false
+				}
+				return nil
+			})
+		}(candidate)
+	}
+	wg.Wait()
+	return
+}
+
+func (broker *MessageQueueBroker) selectBrokers() (brokers []pb.ServerAddress, err error) {
+	candidates, err := broker.selectCandidatesFromMaster(10)
+	if err != nil {
+		return
+	}
+	brokers, err = broker.pickLightestCandidates(candidates, 3)
+	return
+}
+
+func (broker *MessageQueueBroker) selectCandidatesFromMaster(limit int32) (candidates []pb.ServerAddress, err error) {
+	err = broker.withMasterClient(false, broker.MasterClient.GetMaster(), func(client master_pb.SeaweedClient) error {
+		resp, err := client.ListClusterNodes(context.Background(), &master_pb.ListClusterNodesRequest{
+			ClientType: cluster.BrokerType,
+			FilerGroup: broker.option.FilerGroup,
+			Limit:      limit,
+		})
+		if err != nil {
+			return err
+		}
+		if len(resp.ClusterNodes) == 0 {
+			return nil
+		}
+		for _, node := range resp.ClusterNodes {
+			candidates = append(candidates, pb.ServerAddress(node.Address))
+		}
+		return nil
+	})
+	return
+}
+
+type CandidateStatus struct {
+	address      pb.ServerAddress
+	messageCount int64
+	bytesCount   int64
+	load         int64
+}
+
+func (broker *MessageQueueBroker) pickLightestCandidates(candidates []pb.ServerAddress, limit int) (selected []pb.ServerAddress, err error) {
+
+	if len(candidates) <= limit {
+		return candidates, nil
+	}
+
+	candidateStatuses, err := broker.checkBrokerStatus(candidates)
+	if err != nil {
+		return nil, err
+	}
+
+	sort.Slice(candidateStatuses, func(i, j int) bool {
+		return candidateStatuses[i].load < candidateStatuses[j].load
+	})
+
+	for i, candidate := range candidateStatuses {
+		if i >= limit {
+			break
+		}
+		selected = append(selected, candidate.address)
+	}
+
+	return
+}
+
+func (broker *MessageQueueBroker) checkBrokerStatus(candidates []pb.ServerAddress) (candidateStatuses []*CandidateStatus, err error) {
+
+	candidateStatuses = make([]*CandidateStatus, len(candidates))
+	var wg sync.WaitGroup
+	for i, candidate := range candidates {
+		wg.Add(1)
+		go func(i int, candidate pb.ServerAddress) {
+			defer wg.Done()
+			err = broker.withBrokerClient(false, candidate, func(client mq_pb.SeaweedMessagingClient) error {
+				resp, checkErr := client.CheckBrokerLoad(context.Background(), &mq_pb.CheckBrokerLoadRequest{})
+				if checkErr != nil {
+					err = checkErr
+					return err
+				}
+				candidateStatuses[i] = &CandidateStatus{
+					address:      candidate,
+					messageCount: resp.MessageCount,
+					bytesCount:   resp.BytesCount,
+					load:         resp.MessageCount + resp.BytesCount/(64*1024),
+				}
+				return nil
+			})
+		}(i, candidate)
+	}
+	wg.Wait()
+	return
+}
+
+func (broker *MessageQueueBroker) saveSegmentOnFiler(segment *mq.Segment, brokers []pb.ServerAddress) (err error) {
+	return
 }

weed/mq/broker/broker_server.go

@@ -59,3 +59,11 @@ func (broker *MessageQueueBroker) withMasterClient(streamingMode bool, master pb
 	})
 
 }
+
+func (broker *MessageQueueBroker) withBrokerClient(streamingMode bool, server pb.ServerAddress, fn func(client mq_pb.SeaweedMessagingClient) error) error {
+
+	return pb.WithBrokerClient(streamingMode, server, broker.grpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {
+		return fn(client)
+	})
+
+}

weed/mq/topic.go

@@ -1,23 +1,36 @@
 package mq
 
-import "time"
+import (
+	"github.com/chrislusf/seaweedfs/weed/pb/mq_pb"
+	"time"
+)
 
 type Namespace string
 
 type Topic struct {
-	namespace Namespace
-	name      string
+	Namespace Namespace
+	Name      string
 }
 
 type Partition struct {
-	rangeStart int
-	rangeStop  int // exclusive
-	ringSize   int
+	RangeStart int
+	RangeStop  int // exclusive
+	RingSize   int
 }
 
 type Segment struct {
-	topic        Topic
-	id           int32
-	partition    Partition
-	lastModified time.Time
+	Topic        Topic
+	Id           int32
+	Partition    Partition
+	LastModified time.Time
+}
+
+func FromPbSegment(segment *mq_pb.Segment) *Segment {
+	return &Segment{
+		Topic: Topic{
+			Namespace: Namespace(segment.Namespace),
+			Name:      segment.Topic,
+		},
+		Id: segment.Id,
+	}
 }