Add support for deferred publishing via Producer (TCP)

Author: seoester

producer.go

@@ -67,6 +67,7 @@ type Producer struct {
 // producers.
 type ProducerRequest struct {
 	Topic    string
+	Delay    time.Duration
 	Message  []byte
 	Response chan<- error
 	Deadline time.Time
@@ -157,7 +158,7 @@ func (p *Producer) Stop() {
 // 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) Publish(message []byte) (err error) {
+func (p *Producer) Publish(message []byte) error {
 	return p.PublishTo(p.topic, message)
 }
 
@@ -168,7 +169,33 @@ func (p *Producer) Publish(message []byte) (err error) {
 // 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) PublishTo(topic string, message []byte) (err error) {
+func (p *Producer) PublishTo(topic string, message []byte) error {
+	return p.sendProducerRequest(topic, 0, message)
+}
+
+// DeferredPublish sends a deferred message 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) DeferredPublish(delay time.Duration, message []byte) error {
+	return p.DeferredPublishTo(p.topic, delay, message)
+}
+
+// DeferredPublishTo sends a deferred message 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) DeferredPublishTo(topic string, delay time.Duration, message []byte) error {
+	return p.sendProducerRequest(topic, delay, message)
+}
+
+func (p *Producer) sendProducerRequest(topic string, delay time.Duration, message []byte) (err error) {
 	defer func() {
 		if recover() != nil {
 			err = errors.New("publishing to a producer that was already stopped")
@@ -186,6 +213,7 @@ func (p *Producer) PublishTo(topic string, message []byte) (err error) {
 	// it up.
 	p.reqs <- ProducerRequest{
 		Topic:    topic,
+		Delay:    delay,
 		Message:  message,
 		Response: response,
 		Deadline: deadline,
@@ -297,7 +325,7 @@ func (p *Producer) run() {
 				continue
 			}
 
-			if err := p.publish(conn, req.Topic, req.Message); err != nil {
+			if err := p.publishMessage(conn, req.Topic, req.Delay, req.Message); err != nil {
 				req.complete(err)
 				shutdown(err)
 				continue
@@ -365,11 +393,19 @@ 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) publishMessage(conn *Conn, topic string, delay time.Duration, message []byte) error {
+	if delay == 0 {
+		return p.write(conn, Pub{
+			Topic:   topic,
+			Message: message,
+		})
+	} else {
+		return p.write(conn, DPub{
+			Topic:   topic,
+			Delay:   delay,
+			Message: message,
+		})
+	}
 }
 
 func (p *Producer) ping(conn *Conn) error {

producer_test.go

@@ -7,28 +7,67 @@ import (
 	"time"
 )
 
-func TestProducer(t *testing.T) {
-	for _, n := range []int{1, 10, 100, 1000} {
+func mustStartConsumerAndProducer(address, topic string) (*Consumer, *Producer) {
+	c, err := StartConsumer(ConsumerConfig{
+		Topic:   topic,
+		Channel: "channel",
+		Address: address,
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	// Give some time for the consumer to connect.
+	time.Sleep(100 * time.Millisecond)
+
+	p, _ := StartProducer(ProducerConfig{
+		Address:        address,
+		Topic:          topic,
+		MaxConcurrency: 3,
+	})
+	if err != nil {
+		c.Stop()
+		panic(err)
+	}
+
+	return c, p
+}
+
+func consumeAndCheckMessages(c *Consumer, count int, deadline *time.Timer) error {
+	buckets := make([]int, count)
+
+	for i := 0; i != count; i++ {
+		select {
+		case msg := <-c.Messages():
+			b, err := strconv.Atoi(string(msg.Body))
+			if err != nil {
+				return err
+			}
+			buckets[b]++
+			msg.Finish()
+		case <-deadline.C:
+			return fmt.Errorf("timeout")
+		}
+	}
+
+	for i, b := range buckets {
+		if b != 1 {
+			return fmt.Errorf("bucket at index %d has value %d", i, b)
+		}
+	}
+
+	return nil
+}
+
+func TestProducerPublish(t *testing.T) {
+	for _, n := range [...]int{1, 10, 100, 1000} {
 		count := n
-		topic := fmt.Sprintf("test-publisher-%d", n)
+		topic := fmt.Sprintf("test-publish-%d", n)
 		t.Run(topic, func(t *testing.T) {
 			t.Parallel()
 
-			c, _ := StartConsumer(ConsumerConfig{
-				Topic:   topic,
-				Channel: "channel",
-				Address: "localhost:4150",
-			})
+			c, p := mustStartConsumerAndProducer("localhost:4150", topic)
 			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()
 
 			for i := 0; i != count; i++ {
@@ -38,30 +77,60 @@ func TestProducer(t *testing.T) {
 				}
 			}
 
-			buckets := make([]int, count)
-
-			deadline := time.NewTimer(10 * time.Second)
+			// Margin of error: 5*time.Second
+			deadline := time.NewTimer(5 * time.Second)
 			defer deadline.Stop()
 
+			err := consumeAndCheckMessages(c, count, deadline)
+			if err != nil {
+				t.Error(err)
+			}
+		})
+	}
+}
+
+func TestProducerDeferredPublish(t *testing.T) {
+	delay := 10 * time.Second
+
+	for _, n := range [...]int{1, 10, 100, 1000} {
+		count := n
+		topic := fmt.Sprintf("test-deferred-publish-%d", n)
+		t.Run(topic, func(t *testing.T) {
+			t.Parallel()
+
+			c, p := mustStartConsumerAndProducer("localhost:4150", topic)
+			defer c.Stop()
+			defer p.Stop()
+
+			publishStart := time.Now()
+
 			for i := 0; i != count; i++ {
-				select {
-				case msg := <-c.Messages():
-					b, err := strconv.Atoi(string(msg.Body))
-					if err != nil {
-						t.Error(err)
-					}
-					buckets[b]++
-					msg.Finish()
-				case <-deadline.C:
-					t.Error("timeout")
+				if err := p.DeferredPublish(delay, []byte(strconv.Itoa(i))); err != nil {
+					t.Error(err)
 					return
 				}
 			}
 
-			for i, b := range buckets {
-				if b != 1 {
-					t.Errorf("bucket at index %d has value %d", i, b)
-				}
+			publishEnd := time.Now()
+
+			// Margin of error: 1*time.Second
+			delayTimer := time.NewTimer(delay - publishEnd.Sub(publishStart) - 1*time.Second)
+			defer delayTimer.Stop()
+
+			select {
+			case _ = <-c.Messages():
+				t.Error("received deferred message early (before delay time passed)")
+				return
+			case <-delayTimer.C:
+			}
+
+			// Margin of error: 5*time.Second
+			deadline := time.NewTimer(delay - time.Now().Sub(publishEnd) + 5*time.Second)
+			defer deadline.Stop()
+
+			err := consumeAndCheckMessages(c, count, deadline)
+			if err != nil {
+				t.Error(err)
 			}
 		})
 	}