seaweedfs/weed/mq/topic/local_partition.go

package topic

import (
	"context"
	"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/util/log_buffer"
	"google.golang.org/grpc"
	"sync"
	"sync/atomic"
	"time"
)

type LocalPartition struct {
	ListenersWaits   int64
	AckTsNs		  	 int64

	// notifying clients
	ListenersLock sync.Mutex
	ListenersCond *sync.Cond

	Partition
	isLeader        bool
	FollowerBrokers []pb.ServerAddress
	LogBuffer       *log_buffer.LogBuffer
	ConsumerCount   int32
	Publishers      *LocalPartitionPublishers
	Subscribers     *LocalPartitionSubscribers
	FollowerId      int32

	followerStream         mq_pb.SeaweedMessaging_PublishFollowMeClient
	followerGrpcConnection *grpc.ClientConn
	follower               string
}

var TIME_FORMAT = "2006-01-02-15-04-05"

func NewLocalPartition(partition Partition, isLeader bool, followerBrokers []pb.ServerAddress, logFlushFn log_buffer.LogFlushFuncType, readFromDiskFn log_buffer.LogReadFromDiskFuncType) *LocalPartition {
	lp := &LocalPartition{
		Partition:       partition,
		isLeader:        isLeader,
		FollowerBrokers: followerBrokers,
		Publishers:  NewLocalPartitionPublishers(),
		Subscribers: NewLocalPartitionSubscribers(),
	}
	lp.ListenersCond = sync.NewCond(&lp.ListenersLock)
	lp.LogBuffer = log_buffer.NewLogBuffer(fmt.Sprintf("%d/%04d-%04d", partition.UnixTimeNs, partition.RangeStart, partition.RangeStop),
		2*time.Minute, logFlushFn, readFromDiskFn, func() {
			if atomic.LoadInt64(&lp.ListenersWaits) > 0 {
				lp.ListenersCond.Broadcast()
			}
		})
	return lp
}

func (p *LocalPartition) Publish(message *mq_pb.DataMessage) error {
	p.LogBuffer.AddToBuffer(message.Key, message.Value, time.Now().UnixNano())

	// maybe send to the follower
	if p.followerStream != nil {
		println("recv", string(message.Key), message.TsNs)
		if followErr := p.followerStream.Send(&mq_pb.PublishFollowMeRequest{
			Message: &mq_pb.PublishFollowMeRequest_Data{
				Data: message,
			},
		}); followErr != nil {
			return fmt.Errorf("send to follower %s: %v", p.follower, followErr)
		}
	} else {
		atomic.StoreInt64(&p.AckTsNs, message.TsNs)
	}

	return nil
}

func (p *LocalPartition) Subscribe(clientName string, startPosition log_buffer.MessagePosition,
	onNoMessageFn func() bool, eachMessageFn log_buffer.EachLogEntryFuncType) error {
	var processedPosition log_buffer.MessagePosition
	var readPersistedLogErr error
	var readInMemoryLogErr error
	var isDone bool

	for {
		processedPosition, isDone, readPersistedLogErr = p.LogBuffer.ReadFromDiskFn(startPosition, 0, eachMessageFn)
		if readPersistedLogErr != nil {
			glog.V(0).Infof("%s read %v persisted log: %v", clientName, p.Partition, readPersistedLogErr)
			return readPersistedLogErr
		}
		if isDone {
			return nil
		}

		startPosition = processedPosition
		processedPosition, isDone, readInMemoryLogErr = p.LogBuffer.LoopProcessLogData(clientName, startPosition, 0, onNoMessageFn, eachMessageFn)
		if isDone {
			return nil
		}
		startPosition = processedPosition

		if readInMemoryLogErr == log_buffer.ResumeFromDiskError {
			continue
		}
		if readInMemoryLogErr != nil {
			glog.V(0).Infof("%s read %v in memory log: %v", clientName, p.Partition, readInMemoryLogErr)
			return readInMemoryLogErr
		}
	}
}

func (p *LocalPartition) GetEarliestMessageTimeInMemory() time.Time {
	return p.LogBuffer.GetEarliestTime()
}

func (p *LocalPartition) HasData() bool {
	return !p.LogBuffer.GetEarliestTime().IsZero()
}

func (p *LocalPartition) GetEarliestInMemoryMessagePosition() log_buffer.MessagePosition {
	return p.LogBuffer.GetEarliestPosition()
}

func FromPbBrokerPartitionAssignment(self pb.ServerAddress, partition Partition, assignment *mq_pb.BrokerPartitionAssignment, logFlushFn log_buffer.LogFlushFuncType, readFromDiskFn log_buffer.LogReadFromDiskFuncType) *LocalPartition {
	isLeader := assignment.LeaderBroker == string(self)
	followers := make([]pb.ServerAddress, len(assignment.FollowerBrokers))
	for i, followerBroker := range assignment.FollowerBrokers {
		followers[i] = pb.ServerAddress(followerBroker)
	}
	return NewLocalPartition(partition, isLeader, followers, logFlushFn, readFromDiskFn)
}

func (p *LocalPartition) closePublishers() {
	p.Publishers.SignalShutdown()
}
func (p *LocalPartition) closeSubscribers() {
	p.Subscribers.SignalShutdown()
}

func (p *LocalPartition) WaitUntilNoPublishers() {
	for {
		if p.Publishers.IsEmpty() {
			return
		}
		time.Sleep(113 * time.Millisecond)
	}
}

func (p *LocalPartition) MaybeConnectToFollowers(initMessage *mq_pb.PublishMessageRequest_InitMessage, grpcDialOption grpc.DialOption) (err error) {
	if p.followerStream != nil {
		return nil
	}
	if len(initMessage.FollowerBrokers) == 0 {
		return nil
	}

	p.follower = initMessage.FollowerBrokers[0]
	ctx := context.Background()
	p.followerGrpcConnection, err = pb.GrpcDial(ctx, p.follower, true, grpcDialOption)
	if err != nil {
		return fmt.Errorf("fail to dial %s: %v", p.follower, err)
	}
	followerClient := mq_pb.NewSeaweedMessagingClient(p.followerGrpcConnection)
	p.followerStream, err = followerClient.PublishFollowMe(ctx)
	if err != nil {
		return fmt.Errorf("fail to create publish client: %v", err)
	}
	if err = p.followerStream.Send(&mq_pb.PublishFollowMeRequest{
		Message: &mq_pb.PublishFollowMeRequest_Init{
			Init: &mq_pb.PublishFollowMeRequest_InitMessage{
				Topic:      initMessage.Topic,
				Partition:  initMessage.Partition,
			},
		},
	}); err != nil {
		return err
	}

	// start receiving ack from follower
	go func() {
		defer func() {
			println("stop receiving ack from follower")
		}()

		for {
			ack, err := p.followerStream.Recv()
			if err != nil {
				glog.Errorf("Error receiving follower ack: %v", err)
				return
			}
			atomic.StoreInt64(&p.AckTsNs, ack.AckTsNs)
			println("recv ack", ack.AckTsNs)
		}
	}()
	return nil
}

func (p *LocalPartition) MaybeShutdownLocalPartition() (hasShutdown bool) {
	if !p.Publishers.IsEmpty() {
		return
	}
	if !p.Subscribers.IsEmpty() {
		return
	}
	p.LogBuffer.ShutdownLogBuffer()

	if p.followerStream != nil {
		// send close to the follower
		if followErr := p.followerStream.Send(&mq_pb.PublishFollowMeRequest{
			Message: &mq_pb.PublishFollowMeRequest_Close{
				Close: &mq_pb.PublishFollowMeRequest_CloseMessage{},
			},
		}); followErr != nil {
			glog.Errorf("Error closing follower stream: %v", followErr)
		}
		glog.V(4).Infof("closing grpcConnection to follower")
		p.followerGrpcConnection.Close()
		p.followerStream = nil
	}

	return
}

func (p *LocalPartition) Shutdown() {
	p.closePublishers()
	p.closeSubscribers()
	p.LogBuffer.ShutdownLogBuffer()
	atomic.StoreInt32(&p.FollowerId, 0)
	glog.V(0).Infof("local partition %v shutting down", p.Partition)
}