chrislu
2 years ago
9 changed files with 509 additions and 122 deletions
-
33weed/mq/broker/brokder_grpc_pub.go
-
178weed/mq/client/publish_stream_processor.go
-
46weed/mq/client/publisher.go
-
61weed/mq/cmd/qsend/qsend.go
-
4weed/mq/segment/message_serde.go
-
6weed/mq/segment/message_serde_test.go
-
11weed/pb/mq.proto
-
227weed/pb/mq_pb/mq.pb.go
-
39weed/pb/mq_pb/mq_grpc.pb.go
@ -0,0 +1,178 @@ |
|||
package client |
|||
|
|||
import ( |
|||
"context" |
|||
flatbuffers "github.com/google/flatbuffers/go" |
|||
"github.com/seaweedfs/seaweedfs/weed/mq/segment" |
|||
"github.com/seaweedfs/seaweedfs/weed/pb" |
|||
"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb" |
|||
"github.com/seaweedfs/seaweedfs/weed/util" |
|||
"google.golang.org/grpc" |
|||
"google.golang.org/grpc/credentials/insecure" |
|||
"log" |
|||
"sync" |
|||
"sync/atomic" |
|||
"time" |
|||
) |
|||
|
|||
const ( |
|||
batchCountLimit = 3 |
|||
) |
|||
|
|||
type PublishStreamProcessor struct { |
|||
// attributes
|
|||
ProducerId int32 |
|||
ProducerEpoch int32 |
|||
grpcDialOption grpc.DialOption |
|||
|
|||
// input
|
|||
sync.Mutex |
|||
|
|||
timeout time.Duration |
|||
|
|||
// convert into bytes
|
|||
messagesChan chan *Message |
|||
builders chan *flatbuffers.Builder |
|||
batchMessageCountLimit int |
|||
|
|||
messagesSequence int64 |
|||
|
|||
// done channel
|
|||
doneChan chan struct{} |
|||
} |
|||
|
|||
type UploadProcess struct { |
|||
bufferBuilder *flatbuffers.Builder |
|||
batchBuilder *segment.MessageBatchBuilder |
|||
} |
|||
|
|||
func NewPublishStreamProcessor(batchMessageCountLimit int, timeout time.Duration) *PublishStreamProcessor { |
|||
t := &PublishStreamProcessor{ |
|||
grpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()), |
|||
batchMessageCountLimit: batchMessageCountLimit, |
|||
builders: make(chan *flatbuffers.Builder, batchCountLimit), |
|||
messagesChan: make(chan *Message, 1024), |
|||
doneChan: make(chan struct{}), |
|||
timeout: timeout, |
|||
} |
|||
for i := 0; i < batchCountLimit; i++ { |
|||
t.builders <- flatbuffers.NewBuilder(4 * 1024 * 1024) |
|||
} |
|||
go t.doLoopUpload() |
|||
return t |
|||
} |
|||
|
|||
func (p *PublishStreamProcessor) AddMessage(m *Message) error { |
|||
p.messagesChan <- m |
|||
return nil |
|||
} |
|||
|
|||
func (p *PublishStreamProcessor) Shutdown() error { |
|||
p.doneChan <- struct{}{} |
|||
return nil |
|||
} |
|||
|
|||
func (p *PublishStreamProcessor) doFlush(stream mq_pb.SeaweedMessaging_PublishMessageClient, messages []*Message) error { |
|||
|
|||
if len(messages) == 0 { |
|||
return nil |
|||
} |
|||
|
|||
builder := <-p.builders |
|||
bb := segment.NewMessageBatchBuilder(builder, p.ProducerId, p.ProducerEpoch, 3, 4) |
|||
for _, m := range messages { |
|||
bb.AddMessage(p.messagesSequence, m.Ts.UnixNano(), m.Properties, m.Key, m.Content) |
|||
p.messagesSequence++ |
|||
} |
|||
bb.BuildMessageBatch() |
|||
defer func() { |
|||
p.builders <- builder |
|||
}() |
|||
|
|||
return stream.Send(&mq_pb.PublishRequest{ |
|||
Data: &mq_pb.PublishRequest_DataMessage{ |
|||
Message: bb.GetBytes(), |
|||
}, |
|||
}) |
|||
|
|||
} |
|||
|
|||
func (p *PublishStreamProcessor) doLoopUpload() { |
|||
|
|||
brokerGrpcAddress := "localhost:17777" |
|||
|
|||
// TOOD parallelize the uploading with separate uploader
|
|||
messages := make([]*Message, 0, p.batchMessageCountLimit) |
|||
|
|||
util.RetryForever("publish message", func() error { |
|||
return pb.WithBrokerGrpcClient(false, brokerGrpcAddress, p.grpcDialOption, func(client mq_pb.SeaweedMessagingClient) error { |
|||
|
|||
ctx, cancel := context.WithCancel(context.Background()) |
|||
defer cancel() |
|||
|
|||
stream, err := client.PublishMessage(ctx) |
|||
if err != nil { |
|||
log.Printf("grpc PublishMessage: %v", err) |
|||
return err |
|||
} |
|||
|
|||
var atomicStatus int64 |
|||
go func() { |
|||
resp, err := stream.Recv() |
|||
if err != nil { |
|||
log.Printf("response error: %v", err) |
|||
} else { |
|||
log.Printf("response: %v", resp.AckSequence) |
|||
} |
|||
if atomic.LoadInt64(&atomicStatus) < 0 { |
|||
return |
|||
} |
|||
}() |
|||
|
|||
var flushErr error |
|||
// retry previously failed messages
|
|||
if len(messages) >= p.batchMessageCountLimit { |
|||
flushErr = p.doFlush(stream, messages) |
|||
if flushErr != nil { |
|||
return flushErr |
|||
} |
|||
messages = messages[:0] |
|||
} |
|||
|
|||
for { |
|||
select { |
|||
case m := <-p.messagesChan: |
|||
messages = append(messages, m) |
|||
if len(messages) >= p.batchMessageCountLimit { |
|||
if flushErr = p.doFlush(stream, messages); flushErr != nil { |
|||
return flushErr |
|||
} |
|||
messages = messages[:0] |
|||
} |
|||
case <-time.After(p.timeout): |
|||
if flushErr = p.doFlush(stream, messages); flushErr != nil { |
|||
return flushErr |
|||
} |
|||
messages = messages[:0] |
|||
case <-p.doneChan: |
|||
if flushErr = p.doFlush(stream, messages); flushErr != nil { |
|||
return flushErr |
|||
} |
|||
messages = messages[:0] |
|||
println("$ stopping ...") |
|||
break |
|||
} |
|||
} |
|||
|
|||
// stop the response consuming goroutine
|
|||
atomic.StoreInt64(&atomicStatus, -1) |
|||
|
|||
return flushErr |
|||
|
|||
}) |
|||
}, func(err error) (shouldContinue bool) { |
|||
log.Printf("failed with grpc %s: %v", brokerGrpcAddress, err) |
|||
return true |
|||
}) |
|||
|
|||
} |
|||
@ -0,0 +1,46 @@ |
|||
package client |
|||
|
|||
import ( |
|||
"github.com/seaweedfs/seaweedfs/weed/pb" |
|||
"time" |
|||
) |
|||
|
|||
type PublishProcessor interface { |
|||
AddMessage(m *Message) error |
|||
Shutdown() error |
|||
} |
|||
|
|||
type PublisherOption struct { |
|||
Masters string |
|||
Topic string |
|||
} |
|||
|
|||
type Publisher struct { |
|||
option *PublisherOption |
|||
masters []pb.ServerAddress |
|||
processor *PublishStreamProcessor |
|||
} |
|||
|
|||
func NewPublisher(option *PublisherOption) *Publisher { |
|||
p := &Publisher{ |
|||
masters: pb.ServerAddresses(option.Masters).ToAddresses(), |
|||
option: option, |
|||
processor: NewPublishStreamProcessor(3, 887*time.Millisecond), |
|||
} |
|||
return p |
|||
} |
|||
|
|||
type Message struct { |
|||
Key []byte |
|||
Content []byte |
|||
Properties map[string]string |
|||
Ts time.Time |
|||
} |
|||
|
|||
func (p Publisher) Publish(m *Message) error { |
|||
return p.processor.AddMessage(m) |
|||
} |
|||
|
|||
func (p Publisher) Shutdown() error { |
|||
return p.processor.Shutdown() |
|||
} |
|||
@ -0,0 +1,61 @@ |
|||
package main |
|||
|
|||
import ( |
|||
"bufio" |
|||
"flag" |
|||
"fmt" |
|||
"github.com/seaweedfs/seaweedfs/weed/mq/client" |
|||
"os" |
|||
"time" |
|||
) |
|||
|
|||
var ( |
|||
master = flag.String("master", "localhost:9333", "master csv list") |
|||
topic = flag.String("topic", "", "topic name") |
|||
) |
|||
|
|||
func main() { |
|||
flag.Parse() |
|||
|
|||
publisher := client.NewPublisher(&client.PublisherOption{ |
|||
Masters: *master, |
|||
Topic: *topic, |
|||
}) |
|||
|
|||
err := eachLineStdin(func(line string) error { |
|||
if len(line) > 0 { |
|||
if err := publisher.Publish(&client.Message{ |
|||
Key: nil, |
|||
Content: []byte(line), |
|||
Properties: nil, |
|||
Ts: time.Time{}, |
|||
}); err != nil { |
|||
return err |
|||
} |
|||
} |
|||
return nil |
|||
}) |
|||
|
|||
publisher.Shutdown() |
|||
|
|||
if err != nil { |
|||
fmt.Printf("error: %v\n", err) |
|||
} |
|||
} |
|||
|
|||
func eachLineStdin(eachLineFn func(string) error) error { |
|||
scanner := bufio.NewScanner(os.Stdin) |
|||
for scanner.Scan() { |
|||
text := scanner.Text() |
|||
if err := eachLineFn(text); err != nil { |
|||
return err |
|||
} |
|||
} |
|||
|
|||
// handle error
|
|||
if scanner.Err() != nil { |
|||
return fmt.Errorf("scan stdin: %v", scanner.Err()) |
|||
} |
|||
|
|||
return nil |
|||
} |
|||
Write
Preview
Loading…
Cancel
Save
Reference in new issue