support batch

Author: srikanth-mohan

producer.go

@@ -67,7 +67,7 @@ type Producer struct {
 // producers.
 type ProducerRequest struct {
 	Topic    string
-	Message  []byte
+	Messages [][]byte
 	Response chan<- error
 	Deadline time.Time
 }
@@ -158,7 +158,17 @@ func (p *Producer) Stop() {
 // first unsuccessful attempt to publish the message. It is the responsibility
 // of the caller to retry if necessary.
 func (p *Producer) Publish(message []byte) (err error) {
-	return p.PublishTo(p.topic, message)
+	return p.MultiPublishTo(p.topic, [][]byte{message})
+}
+
+// MultiPublish sends message batch using the producer p, returning an error if it was
+// already closed or if an error occurred while publishing the messages.
+//
+// Note that no retry is done internally, the producer will fail after the
+// first unsuccessful attempt to publish the message. It is the responsibility
+// of the caller to retry if necessary.
+func (p *Producer) MultiPublish(message [][]byte) (err error) {
+	return p.MultiPublishTo(p.topic, message)
 }
 
 // PublishTo sends a message to a specific topic using the producer p, returning
@@ -169,6 +179,17 @@ func (p *Producer) Publish(message []byte) (err error) {
 // first unsuccessful attempt to publish the message. It is the responsibility
 // of the caller to retry if necessary.
 func (p *Producer) PublishTo(topic string, message []byte) (err error) {
+	return p.MultiPublishTo(p.topic, [][]byte{message})
+}
+
+// MultiPublishTo sends a message batch to a specific topic using the producer p, returning
+// an error if it was already closed or if an error occurred while publishing the
+// message.
+//
+// Note that no retry is done internally, the producer will fail after the
+// first unsuccessful attempt to publish the message. It is the responsibility
+// of the caller to retry if necessary.
+func (p *Producer) MultiPublishTo(topic string, messages [][]byte) (err error) {
 	defer func() {
 		if recover() != nil {
 			err = errors.New("publishing to a producer that was already stopped")
@@ -186,7 +207,7 @@ func (p *Producer) PublishTo(topic string, message []byte) (err error) {
 	// it up.
 	p.reqs <- ProducerRequest{
 		Topic:    topic,
-		Message:  message,
+		Messages: messages,
 		Response: response,
 		Deadline: deadline,
 	}
@@ -297,7 +318,7 @@ func (p *Producer) run() {
 				continue
 			}
 
-			if err := p.publish(conn, req.Topic, req.Message); err != nil {
+			if err := p.publish(conn, req.Topic, req.Messages); err != nil {
 				req.complete(err)
 				shutdown(err)
 				continue
@@ -365,11 +386,21 @@ func (p *Producer) write(conn *Conn, cmd Command) (err error) {
 	return
 }
 
-func (p *Producer) publish(conn *Conn, topic string, message []byte) error {
-	return p.write(conn, Pub{
-		Topic:   topic,
-		Message: message,
-	})
+func (p *Producer) publish(conn *Conn, topic string, messages [][]byte) error {
+	switch len(messages) {
+	case 0:
+		return nil
+	case 1:
+		return p.write(conn, Pub{
+			Topic:   topic,
+			Message: messages[0],
+		})
+	default:
+		return p.write(conn, MPub{
+			Topic:    topic,
+			Messages: messages,
+		})
+	}
 }
 
 func (p *Producer) ping(conn *Conn) error {

producer_test.go

@@ -66,3 +66,73 @@ func TestProducer(t *testing.T) {
 		})
 	}
 }
+
+func TestProducerBatch(t *testing.T) {
+	batchSize := 10
+	for _, n := range []int{1, 10, 100, 1000} {
+		count := n
+		topic := fmt.Sprintf("test-publisher-batch-%d", n)
+		t.Run(topic, func(t *testing.T) {
+			t.Parallel()
+
+			c, _ := StartConsumer(ConsumerConfig{
+				Topic:       topic,
+				Channel:     "channel",
+				Address:     "localhost:4150",
+				ReadTimeout: 1 * time.Minute,
+			})
+			defer c.Stop()
+
+			// Give some time for the consumer to connect.
+			time.Sleep(100 * time.Millisecond)
+
+			p, _ := StartProducer(ProducerConfig{
+				Address:        "localhost:4150",
+				Topic:          topic,
+				MaxConcurrency: 3,
+			})
+			defer p.Stop()
+			n := 0
+
+			for i := 0; i != count; i++ {
+				batch := make([][]byte, batchSize)
+				for j := 0; j != batchSize; j++ {
+					batch[j] = []byte(strconv.Itoa(n))
+					n++
+				}
+
+				if err := p.MultiPublish(batch); err != nil {
+					t.Error(err)
+					return
+				}
+			}
+
+			buckets := make([]int, count*batchSize)
+
+			deadline := time.NewTimer(10 * time.Second)
+			defer deadline.Stop()
+
+			for i := 0; i != count*batchSize; {
+				select {
+				case msg := <-c.Messages():
+					b, err := strconv.Atoi(string(msg.Body))
+					if err != nil {
+						t.Error(err)
+					}
+					buckets[b]++
+					i++
+					msg.Finish()
+				case <-deadline.C:
+					t.Error("timeout")
+					return
+				}
+			}
+
+			for i, b := range buckets {
+				if b != 1 {
+					t.Errorf("bucket at index %d has value %d", i, b)
+				}
+			}
+		})
+	}
+}