Add support for header transferring (#136)

README.md

@@ -473,6 +473,45 @@ See [`example/cleaner`][cleaner.go].
 
 [cleaner.go]: example/cleaner/main.go
 
+### Header
+
+Redis protocol does not define a specific way to pass additional data like header.
+However, there is often need to pass them (for example for traces propagation).
+
+This implementation injects optional header values marked with a signature into 
+payload body during publishing. When message is consumed, if signature is present, 
+header and original payload are extracted from augmented payload.
+
+Header is defined as `http.Header` for better interoperability with existing libraries,
+for example with [`propagation.HeaderCarrier`](https://pkg.go.dev/go.opentelemetry.io/otel/propagation#HeaderCarrier).
+
+```go
+ // ....
+
+ h := make(http.Header)
+ h.Set("X-Baz", "quux")
+
+ // You can add header to your payload during publish.
+ _ = pub.Publish(rmq.PayloadWithHeader(`{"foo":"bar"}`, h))
+
+ // ....
+
+ _, _ = con.AddConsumerFunc("tag", func(delivery rmq.Delivery) {
+     // And receive header back in consumer.
+     delivery.(rmq.WithHeader).Header().Get("X-Baz") // "quux"
+
+     // ....
+ })
+```
+
+Adding a header is an explicit opt-in operation and so it does not affect library's
+backwards compatibility by default (when not used). 
+
+Please note that adding header may lead to compatibility issues if:
+* consumer is built with older version of `rmq` when publisher has already 
+   started using header, this can be avoided by upgrading consumers before publishers;
+* consumer is not using `rmq` (other libs, low level tools like `redis-cli`) and is 
+   not aware of payload format extension.
 
 ## Testing Included
 

delivery.go

@@ -3,6 +3,7 @@ package rmq
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"time"
 )
 
@@ -14,16 +15,26 @@ type Delivery interface {
 	Push() error
 }
 
+var (
+	_ Delivery   = &redisDelivery{}
+	_ WithHeader = &redisDelivery{}
+)
+
 type redisDelivery struct {
 	ctx         context.Context
 	payload     string
+	header      http.Header
 	unackedKey  string
 	rejectedKey string
 	pushKey     string
 	redisClient RedisClient
 	errChan     chan<- error
 }
 
+func (delivery *redisDelivery) Header() http.Header {
+	return delivery.header
+}
+
 func newDelivery(
 	ctx context.Context,
 	payload string,
@@ -32,8 +43,8 @@ func newDelivery(
 	pushKey string,
 	redisClient RedisClient,
 	errChan chan<- error,
-) *redisDelivery {
-	return &redisDelivery{
+) (*redisDelivery, error) {
+	rd := redisDelivery{
 		ctx:         ctx,
 		payload:     payload,
 		unackedKey:  unackedKey,
@@ -42,6 +53,14 @@ func newDelivery(
 		redisClient: redisClient,
 		errChan:     errChan,
 	}
+
+	var err error
+
+	if rd.header, rd.payload, err = ExtractHeaderAndPayload(payload); err != nil {
+		return nil, err
+	}
+
+	return &rd, nil
 }
 
 func (delivery *redisDelivery) String() string {

header.go

@@ -0,0 +1,79 @@
+package rmq
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"strings"
+)
+
+// Redis protocol does not define a specific way to pass additional data like header.
+// However, there is often need to pass them (for example for traces propagation).
+//
+// This implementation injects optional header values marked with a signature into payload body
+// during publishing. When message is consumed, if signature is present, header and original payload
+// are extracted from augmented payload.
+//
+// Header is defined as http.Header for better interoperability with existing libraries,
+// for example with go.opentelemetry.io/otel/propagation.HeaderCarrier.
+
+// PayloadWithHeader creates a payload string with header.
+func PayloadWithHeader(payload string, header http.Header) string {
+	if len(header) == 0 {
+		return payload
+	}
+
+	hd, _ := json.Marshal(header) // String map never fails marshaling.
+
+	return jsonHeaderSignature + string(hd) + "\n" + payload
+}
+
+// PayloadBytesWithHeader creates payload bytes slice with header.
+func PayloadBytesWithHeader(payload []byte, header http.Header) []byte {
+	if len(header) == 0 {
+		return payload
+	}
+
+	hd, _ := json.Marshal(header) // String map never fails marshaling.
+
+	res := make([]byte, 0, len(jsonHeaderSignature)+len(hd)+1+len(payload))
+	res = append(res, []byte(jsonHeaderSignature)...)
+	res = append(res, hd...)
+	res = append(res, '\n')
+	res = append(res, payload...)
+
+	return res
+}
+
+// ExtractHeaderAndPayload splits augmented payload into header and original payload if specific signature is present.
+func ExtractHeaderAndPayload(payload string) (http.Header, string, error) {
+	if !strings.HasPrefix(payload, jsonHeaderSignature) {
+		return nil, payload, nil
+	}
+
+	lineEnd := strings.Index(payload, "\n")
+	if lineEnd == -1 {
+		return nil, "", errors.New("missing line separator")
+	}
+
+	first := payload[len(jsonHeaderSignature):lineEnd]
+	rest := payload[lineEnd+1:]
+
+	header := make(http.Header)
+
+	if err := json.Unmarshal([]byte(first), &header); err != nil {
+		return nil, "", fmt.Errorf("parsing header: %w", err)
+	}
+
+	return header, rest, nil
+}
+
+// WithHeader is a Delivery with Header.
+type WithHeader interface {
+	Header() http.Header
+}
+
+// jsonHeaderSignature is a signature marker to indicate JSON header presence.
+// Do not change the value.
+const jsonHeaderSignature = "\xFF\x00\xBE\xBEJ"

header_test.go

@@ -0,0 +1,112 @@
+package rmq_test
+
+import (
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/adjust/rmq/v5"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPayloadWithHeader(t *testing.T) {
+	p := `{"foo":"bar"}`
+
+	h := make(http.Header)
+	ph := rmq.PayloadWithHeader(p, h)
+	assert.Equal(t, p, ph) // No change for empty header.
+	h2, p2, err := rmq.ExtractHeaderAndPayload(ph)
+	require.NoError(t, err)
+	assert.Nil(t, h2)
+	assert.Equal(t, p, p2)
+
+	h.Set("X-Foo", "Bar")
+	ph = rmq.PayloadWithHeader(p, h)
+	assert.NotEqual(t, p, ph)
+
+	h2, p2, err = rmq.ExtractHeaderAndPayload(ph)
+	require.NoError(t, err)
+	assert.Equal(t, h, h2)
+	assert.Equal(t, p, p2)
+}
+
+func TestPayloadBytesWithHeader(t *testing.T) {
+	p := `{"foo":"bar"}`
+
+	h := make(http.Header)
+	ph := rmq.PayloadBytesWithHeader([]byte(p), h)
+	assert.Equal(t, p, string(ph)) // No change for empty header.
+	h2, p2, err := rmq.ExtractHeaderAndPayload(string(ph))
+	require.NoError(t, err)
+	assert.Nil(t, h2)
+	assert.Equal(t, p, p2)
+
+	h.Set("X-Foo", "Bar")
+	ph = rmq.PayloadBytesWithHeader([]byte(p), h)
+	assert.NotEqual(t, p, ph)
+
+	h2, p2, err = rmq.ExtractHeaderAndPayload(string(ph))
+	require.NoError(t, err)
+	assert.Equal(t, h, h2)
+	assert.Equal(t, p, string(p2))
+}
+
+func TestExtractHeaderAndPayload(t *testing.T) {
+	t.Run("missing_line_separator", func(t *testing.T) {
+		ph := rmq.PayloadWithHeader("foo", http.Header{"foo": []string{"bar"}})
+		ph = ph[0:7] // Truncating payload.
+		h, p, err := rmq.ExtractHeaderAndPayload(ph)
+		require.Error(t, err)
+		assert.Nil(t, h)
+		assert.Empty(t, p)
+	})
+
+	t.Run("invalid_json", func(t *testing.T) {
+		ph := rmq.PayloadWithHeader("foo", http.Header{"foo": []string{"bar"}})
+		ph = strings.Replace(ph, `"`, `'`, 1) // Corrupting JSON.
+		h, p, err := rmq.ExtractHeaderAndPayload(ph)
+		require.Error(t, err)
+		assert.Nil(t, h)
+		assert.Empty(t, p)
+	})
+
+	t.Run("ok", func(t *testing.T) {
+		ph := rmq.PayloadWithHeader("foo", http.Header{"foo": []string{"bar"}})
+		h, p, err := rmq.ExtractHeaderAndPayload(ph)
+		require.NoError(t, err)
+		assert.Equal(t, http.Header{"foo": []string{"bar"}}, h)
+		assert.Equal(t, "foo", p)
+	})
+
+	t.Run("ok_line_breaks", func(t *testing.T) {
+		ph := rmq.PayloadWithHeader("foo", http.Header{"foo": []string{"bar1\nbar2\nbar3"}})
+		h, p, err := rmq.ExtractHeaderAndPayload(ph)
+		require.NoError(t, err)
+		assert.Equal(t, http.Header{"foo": []string{"bar1\nbar2\nbar3"}}, h)
+		assert.Equal(t, "foo", p)
+	})
+}
+
+func ExamplePayloadWithHeader() {
+	var (
+		pub, con rmq.Queue
+	)
+
+	// ....
+
+	h := make(http.Header)
+	h.Set("X-Baz", "quux")
+
+	// You can add header to your payload during publish.
+	_ = pub.Publish(rmq.PayloadWithHeader(`{"foo":"bar"}`, h))
+
+	// ....
+
+	_, _ = con.AddConsumerFunc("tag", func(delivery rmq.Delivery) {
+		// And receive header back in consumer.
+		delivery.(rmq.WithHeader).Header().Get("X-Baz") // "quux"
+
+		// ....
+	})
+}

queue.go

@@ -216,13 +216,18 @@ func (queue *redisQueue) consumeBatch() error {
 			return err
 		}
 
-		queue.deliveryChan <- queue.newDelivery(payload)
+		d, err := queue.newDelivery(payload)
+		if err != nil {
+			return err
+		}
+
+		queue.deliveryChan <- d
 	}
 
 	return nil
 }
 
-func (queue *redisQueue) newDelivery(payload string) Delivery {
+func (queue *redisQueue) newDelivery(payload string) (Delivery, error) {
 	return newDelivery(
 		queue.ackCtx,
 		payload,

queue_test.go

@@ -3,7 +3,9 @@ package rmq
 import (
 	"fmt"
 	"math"
+	"net/http"
 	"strconv"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -833,3 +835,54 @@ func TestQueueDrain(t *testing.T) {
 		eventuallyReady(t, queue, int64(100-x*10))
 	}
 }
+
+func TestQueueHeader(t *testing.T) {
+	redisAddr, closer := testRedis(t)
+	defer closer()
+
+	connection, err := OpenConnection("queue-h-conn", "tcp", redisAddr, 1, nil)
+	assert.NoError(t, err)
+	require.NotNil(t, connection)
+
+	queue, err := connection.OpenQueue("queue-h")
+	assert.NoError(t, err)
+	require.NotNil(t, queue)
+	_, err = queue.PurgeReady()
+	assert.NoError(t, err)
+
+	assert.NoError(t, queue.StartConsuming(2, time.Millisecond))
+	time.Sleep(time.Millisecond)
+	assert.NoError(t, err)
+
+	consumed := int64(0)
+	cons := ConsumerFunc(func(delivery Delivery) {
+		atomic.AddInt64(&consumed, 1)
+
+		h, ok := delivery.(WithHeader)
+		assert.True(t, ok)
+
+		switch delivery.Payload() {
+		case "queue-d1":
+			assert.Empty(t, h.Header())
+		case "queue-d2":
+			require.NotNil(t, h.Header())
+			assert.Equal(t, "d2", h.Header().Get("X-Foo"))
+		default:
+			assert.Failf(t, "unexpected payload: %q", delivery.Payload())
+		}
+
+	})
+
+	_, err = queue.AddConsumer("queue-cons1", cons)
+	assert.NoError(t, err)
+
+	assert.NoError(t, queue.Publish("queue-d1"))
+	assert.NoError(t, queue.Publish(PayloadWithHeader("queue-d2", http.Header{"X-Foo": []string{"d2"}})))
+
+	assert.Eventually(t, func() bool {
+		return atomic.LoadInt64(&consumed) == 2
+	}, 10*time.Second, time.Millisecond)
+
+	<-queue.StopConsuming()
+	assert.NoError(t, connection.stopHeartbeat())
+}