Merge pull request #54 from heetch/rog-002-sendmessages

producer: add SendMessages method

Author: rogpeppe

producer/producer.go

@@ -2,6 +2,7 @@ package producer
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/Shopify/sarama"
 	"github.com/heetch/felice/codec"
@@ -65,14 +66,70 @@ func (p *Producer) Send(ctx context.Context, topic string, body interface{}, opt
 	return msg, err
 }
 
+// SendMessagesErrors is the error type returned if SendMessages
+// fails to send to Kafka.
+type SendMessagesErrors []*SendMessagesError
+
+func (e SendMessagesErrors) Error() string {
+	return fmt.Sprintf("kafka: failed to deliver %d messages", len(e))
+}
+
+// SendMessagesError describes why one message
+// failed to be sent.
+type SendMessagesError struct {
+	Msg *Message
+	Err error
+}
+
+// SendMessages sends all the given messages in order.
+// If it fails to send the messages to Kafka, it will return a
+// SendMessagesErrors error describing which messages failed.
+func (p *Producer) SendMessages(ctx context.Context, msgs []*Message) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+	if p.config.Converter == nil {
+		return errors.New("producer: missing Converter in config")
+	}
+	pmsgs := make([]*sarama.ProducerMessage, len(msgs))
+	for i, m := range msgs {
+		m.prepare()
+		pmsg, err := p.config.Converter.ToKafka(ctx, m)
+		if err != nil {
+			return err
+		}
+		pmsg.Metadata = m
+		pmsgs[i] = pmsg
+	}
+	err := p.SyncProducer.SendMessages(pmsgs)
+	// Some messages may still have been sent, so copy all the
+	// send information across even when there are errors.
+	for i, m := range msgs {
+		p := pmsgs[i]
+		m.Partition, m.Offset, m.ProducedAt = p.Partition, p.Offset, p.Timestamp
+	}
+	if producerErrs, ok := err.(sarama.ProducerErrors); ok {
+		sendErrs := make(SendMessagesErrors, len(producerErrs))
+		for i, e := range producerErrs {
+			sendErrs[i] = &SendMessagesError{
+				Msg: e.Msg.Metadata.(*Message),
+				Err: e.Err,
+			}
+		}
+		err = sendErrs
+	}
+	return err
+}
+
 // SendMessage sends the given message to Kafka synchronously.
 func (p *Producer) SendMessage(ctx context.Context, msg *Message) error {
 	select {
 	case <-ctx.Done():
 		return ctx.Err()
 	default:
 	}
-
 	if p.config.Converter == nil {
 		return errors.New("producer: missing Converter in config")
 	}

producer/producer_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/Shopify/sarama"
 	"github.com/Shopify/sarama/mocks"
 	"github.com/heetch/felice/producer"
 	"github.com/stretchr/testify/require"
@@ -39,6 +40,73 @@ func TestSendMessage(t *testing.T) {
 	}
 }
 
+func TestSendMessages(t *testing.T) {
+	msgs := []*producer.Message{
+		producer.NewMessage("topic1", "message1"),
+		&producer.Message{Topic: "topic2", Body: "message2"},
+		&producer.Message{Topic: "topic2", Body: "message3"},
+	}
+	cfg := producer.NewConfig("id", producer.MessageConverterV1())
+	msp := mocks.NewSyncProducer(t, &cfg.Config)
+	p, err := producer.NewFrom(msp, cfg)
+	require.NoError(t, err)
+
+	for i, msg := range msgs {
+		i, msg := i, msg
+		msp.ExpectSendMessageWithCheckerFunctionAndSucceed(func(val []byte) error {
+			if got, want := string(val), fmt.Sprintf("%q", msg.Body); got != want {
+				return fmt.Errorf("unexpected message %d; got %q want %q", i, got, want)
+			}
+			return nil
+		})
+	}
+	err = p.SendMessages(context.Background(), msgs)
+	require.NoError(t, err)
+
+	// Unfortunately the Sarama mock doesn't fill out of any of the Offset, Partition
+	// or Timestamp fields when SendMessages is called, so we can't test
+	// that functionality here.
+}
+
+func TestSendMessagesError(t *testing.T) {
+	msgs := []*producer.Message{
+		producer.NewMessage("topic1", "message1"),
+		&producer.Message{Topic: "topic2", Body: "message2"},
+		&producer.Message{Topic: "topic2", Body: "message3"},
+	}
+	perr1 := fmt.Errorf("first error")
+	perr2 := fmt.Errorf("second error")
+	cfg := producer.NewConfig("id", producer.MessageConverterV1())
+	sendMessages := func(smsgs []*sarama.ProducerMessage) error {
+		return sarama.ProducerErrors{{
+			Msg: smsgs[0],
+			Err: perr1,
+		}, {
+			Msg: smsgs[2],
+			Err: perr2,
+		}}
+	}
+	p, err := producer.NewFrom(sendMessagesFunc{f: sendMessages}, cfg)
+	require.NoError(t, err)
+	err = p.SendMessages(context.Background(), msgs)
+	sendErrs, ok := err.(producer.SendMessagesErrors)
+	require.True(t, ok)
+	require.Len(t, sendErrs, 2)
+
+	if sendErrs[0].Msg != msgs[0] {
+		t.Errorf("unexpected message at 0: %#v", sendErrs[0].Msg)
+	}
+	if sendErrs[0].Err != perr1 {
+		t.Errorf("unexpected error at 0: %#v", sendErrs[0].Err)
+	}
+	if sendErrs[1].Msg != msgs[2] {
+		t.Errorf("unexpected message at 0: %#v", sendErrs[0].Msg)
+	}
+	if sendErrs[1].Err != perr2 {
+		t.Errorf("unexpected error at 0: %#v", sendErrs[0].Err)
+	}
+}
+
 func TestSend(t *testing.T) {
 	msp := mocks.NewSyncProducer(t, nil)
 	cfg := producer.NewConfig("id", producer.MessageConverterV1())
@@ -58,3 +126,12 @@ func TestSend(t *testing.T) {
 	_, err = p.Send(context.Background(), "topic", "message")
 	require.EqualError(t, err, "producer: failed to send message: cannot produce message")
 }
+
+type sendMessagesFunc struct {
+	sarama.SyncProducer
+	f func(msgs []*sarama.ProducerMessage) error
+}
+
+func (f sendMessagesFunc) SendMessages(msgs []*sarama.ProducerMessage) error {
+	return f.f(msgs)
+}