Merge pull request #57 from segmentio/support-batch

Support publishing batch messages

Author: srikanth-mohan

.github/workflows/ci.yml

@@ -8,15 +8,15 @@ jobs:
       # github.com/actions/setup-go/tags
       uses: actions/setup-go@v5
       with:
-        go-version: 1.22.x
+        go-version: 1.24.x
       # github.com/actions/checkout/tags
     - uses: actions/checkout@v4
       with:
          path: './src/github.com/segmentio/nsq-go'
     - run: echo "PATH=$GITHUB_WORKSPACE/bin:$PATH" >> $GITHUB_ENV
     - name: Run tests
       run: |
-        docker-compose up -d
+        docker compose up -d
         go vet ./...
         go run honnef.co/go/tools/cmd/staticcheck@latest ./...
         go test -race -v ./...

docker-compose.yml

@@ -1,5 +1,3 @@
-version: '2'
-
 # Use this docker-compose file to setup the test environment before running the
 # tests.
 services:

go.mod

@@ -1,6 +1,6 @@
 module github.com/segmentio/nsq-go
 
-go 1.18
+go 1.24
 
 require (
 	github.com/pkg/errors v0.9.1

go.sum

@@ -1,5 +1,7 @@
 github.com/kr/pretty v0.2.1 h1:Fmg33tUaq4/8ym9TJN1x7sLJnHVwhP33CNkpYV/7rwI=
+github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
+github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc=
@@ -21,7 +23,9 @@ golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
 golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
 gopkg.in/go-playground/assert.v1 v1.2.1 h1:xoYuJVE7KT85PYWrN730RguIQO0ePzVRfFMXadIrXTM=
+gopkg.in/go-playground/assert.v1 v1.2.1/go.mod h1:9RXL0bg/zibRAgZUYszZSwO/z8Y/a8bDuhia5mkpMnE=
 gopkg.in/go-playground/mold.v2 v2.2.0 h1:Y4IYB4/HYQfuq43zaKh6vs9cVelLE9qbqe2fkyfCTWQ=
 gopkg.in/go-playground/mold.v2 v2.2.0/go.mod h1:XMyyRsGtakkDPbxXbrA5VODo6bUXyvoDjLd5l3T0XoA=
 gopkg.in/validator.v2 v2.0.1 h1:xF0KWyGWXm/LM2G1TrEjqOu4pa6coO9AlWSf3msVfDY=

nsqlookup/engine_test.go

@@ -15,10 +15,6 @@ const (
 	tombTimeout = 50 * time.Millisecond
 )
 
-func init() {
-	rand.Seed(time.Now().UnixNano())
-}
-
 func testEngine(t *testing.T, do func(context.Context, *testing.T, Engine)) {
 	tests := []struct {
 		Type string

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)
+				}
+			}
+		})
+	}
+}