gitlab-com-seaweedfs/seaweedfs
1
0
Fork 0
mirror of https://github.com/seaweedfs/seaweedfs.git synced 2024-01-19 02:48:24 +00:00
Code Issues Projects Releases Wiki Activity

subscribe support cancel

Browse source
This commit is contained in:
Chris Lu 2020-05-17 11:10:45 -07:00
parent f11233cd49
commit 95ca9dd8a2
2 changed files with 26 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
weed/messaging/msgclient/sub_chan.go
View file

@ -1,6 +1,7 @@
package msgclient package msgclient
import ( import (
"context"
"crypto/md5" "crypto/md5"
"hash" "hash"
"io" "io"
@ -15,6 +16,7 @@ type SubChannel struct {
ch chan []byte ch chan []byte
stream messaging_pb.SeaweedMessaging_SubscribeClient stream messaging_pb.SeaweedMessaging_SubscribeClient
md5hash hash.Hash md5hash hash.Hash
cancel context.CancelFunc
} }
func (mc *MessagingClient) NewSubChannel(subscriberId, chanName string) (*SubChannel, error) { func (mc *MessagingClient) NewSubChannel(subscriberId, chanName string) (*SubChannel, error) {
@ -27,7 +29,8 @@ func (mc *MessagingClient) NewSubChannel(subscriberId, chanName string) (*SubCha
if err != nil { if err != nil {
return nil, err return nil, err
} }
sc, err := setupSubscriberClient(grpcConnection, tp, subscriberId, time.Unix(0, 0)) ctx, cancel := context.WithCancel(context.Background())
sc, err := setupSubscriberClient(ctx, grpcConnection, tp, subscriberId, time.Unix(0, 0))
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -36,6 +39,7 @@ func (mc *MessagingClient) NewSubChannel(subscriberId, chanName string) (*SubCha
ch: make(chan []byte), ch: make(chan []byte),
stream: sc, stream: sc,
md5hash: md5.New(), md5hash: md5.New(),
cancel: cancel,
} }
go func() { go func() {
@ -57,6 +61,7 @@ func (mc *MessagingClient) NewSubChannel(subscriberId, chanName string) (*SubCha
IsClose: true, IsClose: true,
}) })
close(t.ch) close(t.ch)
cancel()
return return
} }
t.ch <- resp.Data.Value t.ch <- resp.Data.Value
@ -74,3 +79,7 @@ func (sc *SubChannel) Channel() chan []byte {
func (sc *SubChannel) Md5() []byte { func (sc *SubChannel) Md5() []byte {
return sc.md5hash.Sum(nil) return sc.md5hash.Sum(nil)
} }
func (sc *SubChannel) Cancel() {
sc.cancel()
}

19
weed/messaging/msgclient/subscriber.go
View file

@ -12,6 +12,7 @@ import (
type Subscriber struct { type Subscriber struct {
subscriberClients []messaging_pb.SeaweedMessaging_SubscribeClient subscriberClients []messaging_pb.SeaweedMessaging_SubscribeClient
subscriberCancels []context.CancelFunc
subscriberId string subscriberId string
} }
@ -21,6 +22,7 @@ func (mc *MessagingClient) NewSubscriber(subscriberId, namespace, topic string,
PartitionCount: 4, PartitionCount: 4,
} }
subscriberClients := make([]messaging_pb.SeaweedMessaging_SubscribeClient, topicConfiguration.PartitionCount) subscriberClients := make([]messaging_pb.SeaweedMessaging_SubscribeClient, topicConfiguration.PartitionCount)
subscriberCancels := make([]context.CancelFunc, topicConfiguration.PartitionCount)
for i := 0; i < int(topicConfiguration.PartitionCount); i++ { for i := 0; i < int(topicConfiguration.PartitionCount); i++ {
if partitionId>=0 && i != partitionId { if partitionId>=0 && i != partitionId {
@ -35,21 +37,24 @@ func (mc *MessagingClient) NewSubscriber(subscriberId, namespace, topic string,
if err != nil { if err != nil {
return nil, err return nil, err
} }
client, err := setupSubscriberClient(grpcClientConn, tp, subscriberId, startTime) ctx, cancel := context.WithCancel(context.Background())
client, err := setupSubscriberClient(ctx, grpcClientConn, tp, subscriberId, startTime)
if err != nil { if err != nil {
return nil, err return nil, err
} }
subscriberClients[i] = client subscriberClients[i] = client
subscriberCancels[i] = cancel
} }
return &Subscriber{ return &Subscriber{
subscriberClients: subscriberClients, subscriberClients: subscriberClients,
subscriberCancels: subscriberCancels,
subscriberId: subscriberId, subscriberId: subscriberId,
}, nil }, nil
} }
func setupSubscriberClient(grpcConnection *grpc.ClientConn, tp broker.TopicPartition, subscriberId string, startTime time.Time) (stream messaging_pb.SeaweedMessaging_SubscribeClient, err error) { func setupSubscriberClient(ctx context.Context, grpcConnection *grpc.ClientConn, tp broker.TopicPartition, subscriberId string, startTime time.Time) (stream messaging_pb.SeaweedMessaging_SubscribeClient, err error) {
stream, err = messaging_pb.NewSeaweedMessagingClient(grpcConnection).Subscribe(context.Background()) stream, err = messaging_pb.NewSeaweedMessagingClient(grpcConnection).Subscribe(ctx)
if err != nil { if err != nil {
return return
} }
@ -98,3 +103,11 @@ func (s *Subscriber) Subscribe(processFn func(m *messaging_pb.Message)) {
} }
} }
} }
func (s *Subscriber) Shutdown() {
for i := 0; i < len(s.subscriberClients); i++ {
if s.subscriberCancels[i] != nil {
s.subscriberCancels[i]()
}
}
}