Support ordering key

pkg/googlecloud/marshaler.go

@@ -30,7 +30,7 @@ type MarshalerUnmarshaler interface {
 	Unmarshaler
 }
 
-func (m DefaultMarshalerUnmarshaler) Marshal(topic string, msg *message.Message) (*pubsub.Message, error) {
+func (DefaultMarshalerUnmarshaler) Marshal(topic string, msg *message.Message) (*pubsub.Message, error) {
 	if value := msg.Metadata.Get(UUIDHeaderKey); value != "" {
 		return nil, errors.Errorf("metadata %s is reserved by watermill for message UUID", UUIDHeaderKey)
 	}
@@ -51,7 +51,7 @@ func (m DefaultMarshalerUnmarshaler) Marshal(topic string, msg *message.Message)
 	return marshaledMsg, nil
 }
 
-func (u DefaultMarshalerUnmarshaler) Unmarshal(pubsubMsg *pubsub.Message) (*message.Message, error) {
+func (DefaultMarshalerUnmarshaler) Unmarshal(pubsubMsg *pubsub.Message) (*message.Message, error) {
 	metadata := make(message.Metadata, len(pubsubMsg.Attributes))
 
 	var id string
@@ -70,3 +70,61 @@ func (u DefaultMarshalerUnmarshaler) Unmarshal(pubsubMsg *pubsub.Message) (*mess
 
 	return msg, nil
 }
+
+type GenerateOrderingKey func(topic string, msg *message.Message) (string, error)
+
+type orderingMarshaler struct {
+	Marshaler
+
+	generateOrderingKey GenerateOrderingKey
+}
+
+func NewOrderingMarshaler(generateOrderingKey GenerateOrderingKey) Marshaler {
+	return &orderingMarshaler{
+		Marshaler:           DefaultMarshalerUnmarshaler{},
+		generateOrderingKey: generateOrderingKey,
+	}
+}
+
+func (om orderingMarshaler) Marshal(topic string, msg *message.Message) (*pubsub.Message, error) {
+	marshaledMsg, err := om.Marshaler.Marshal(topic, msg)
+	if err != nil {
+		return nil, err
+	}
+
+	orderingKey, err := om.generateOrderingKey(topic, msg)
+	if err != nil {
+		return nil, errors.Wrap(err, "cannot generate ordering key")
+	}
+	marshaledMsg.OrderingKey = orderingKey
+
+	return marshaledMsg, nil
+}
+
+type ExtractOrderingKey func(orderingKey string, msg *message.Message) error
+
+type orderingUnmarshaler struct {
+	Unmarshaler
+
+	extractOrderingKey ExtractOrderingKey
+}
+
+func NewOrderingUnmarshaler(extractOrderingKey ExtractOrderingKey) Unmarshaler {
+	return &orderingUnmarshaler{
+		Unmarshaler:        DefaultMarshalerUnmarshaler{},
+		extractOrderingKey: extractOrderingKey,
+	}
+}
+
+func (ou orderingUnmarshaler) Unmarshal(pubsubMsg *pubsub.Message) (*message.Message, error) {
+	msg, err := ou.Unmarshaler.Unmarshal(pubsubMsg)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := ou.extractOrderingKey(pubsubMsg.OrderingKey, msg); err != nil {
+		return nil, errors.Wrap(err, "cannot extract ordering key")
+	}
+
+	return msg, nil
+}

pkg/googlecloud/publisher.go

@@ -38,6 +38,8 @@ type PublisherConfig struct {
 	// If false (default), `Publisher` tries to create a topic if there is none with the requested name.
 	// Otherwise, trying to subscribe to non-existent subscription results in `ErrTopicDoesNotExist`.
 	DoNotCreateTopicIfMissing bool
+	// Enables the topic message ordering
+	EnableMessageOrdering bool
 
 	// ConnectTimeout defines the timeout for connecting to Pub/Sub
 	ConnectTimeout time.Duration
@@ -170,6 +172,7 @@ func (p *Publisher) topic(ctx context.Context, topic string) (t *pubsub.Topic, e
 	}()
 
 	t = p.client.Topic(topic)
+	t.EnableMessageOrdering = p.config.EnableMessageOrdering
 
 	// todo: theoretically, one could want different publish settings per topic, which is supported by the client lib
 	// different instances of publisher may be used then

pkg/googlecloud/pubsub_test.go

@@ -19,12 +19,13 @@ import (
 
 // Run `docker-compose up` and set PUBSUB_EMULATOR_HOST=localhost:8085 for this to work
 
-func newPubSub(t *testing.T, marshaler googlecloud.MarshalerUnmarshaler, subscriptionName googlecloud.SubscriptionNameFn) (message.Publisher, message.Subscriber) {
+func newPubSub(t *testing.T, enableMessageOrdering bool, marshaler googlecloud.Marshaler, unmarshaler googlecloud.Unmarshaler, subscriptionName googlecloud.SubscriptionNameFn) (message.Publisher, message.Subscriber) {
 	logger := watermill.NewStdLogger(true, true)
 
 	publisher, err := googlecloud.NewPublisher(
 		googlecloud.PublisherConfig{
-			Marshaler: marshaler,
+			EnableMessageOrdering: enableMessageOrdering,
+			Marshaler:             marshaler,
 		},
 		logger,
 	)
@@ -34,9 +35,10 @@ func newPubSub(t *testing.T, marshaler googlecloud.MarshalerUnmarshaler, subscri
 		googlecloud.SubscriberConfig{
 			GenerateSubscriptionName: subscriptionName,
 			SubscriptionConfig: pubsub.SubscriptionConfig{
-				RetainAckedMessages: false,
+				RetainAckedMessages:   false,
+				EnableMessageOrdering: enableMessageOrdering,
 			},
-			Unmarshaler: marshaler,
+			Unmarshaler: unmarshaler,
 		},
 		logger,
 	)
@@ -45,14 +47,44 @@ func newPubSub(t *testing.T, marshaler googlecloud.MarshalerUnmarshaler, subscri
 	return publisher, subscriber
 }
 
+func createPubSub(t *testing.T) (message.Publisher, message.Subscriber) {
+	var defaultMarshalerUnmarshaler googlecloud.DefaultMarshalerUnmarshaler
+	return newPubSub(t, false, defaultMarshalerUnmarshaler, defaultMarshalerUnmarshaler, googlecloud.TopicSubscriptionName)
+}
+
 func createPubSubWithSubscriptionName(t *testing.T, subscriptionName string) (message.Publisher, message.Subscriber) {
-	return newPubSub(t, googlecloud.DefaultMarshalerUnmarshaler{},
+	var defaultMarshalerUnmarshaler googlecloud.DefaultMarshalerUnmarshaler
+	return newPubSub(t, false, defaultMarshalerUnmarshaler, defaultMarshalerUnmarshaler,
 		googlecloud.TopicSubscriptionNameWithSuffix(subscriptionName),
 	)
 }
 
-func createPubSub(t *testing.T) (message.Publisher, message.Subscriber) {
-	return newPubSub(t, googlecloud.DefaultMarshalerUnmarshaler{}, googlecloud.TopicSubscriptionName)
+func createPubSubWithOrdering(t *testing.T) (message.Publisher, message.Subscriber) {
+	return newPubSub(
+		t,
+		true,
+		googlecloud.NewOrderingMarshaler(func(topic string, msg *message.Message) (string, error) {
+			return "ordering_key", nil
+		}),
+		googlecloud.NewOrderingUnmarshaler(func(orderingKey string, msg *message.Message) error {
+			return nil
+		}),
+		googlecloud.TopicSubscriptionName,
+	)
+}
+
+func createPubSubWithSubscriptionNameWithOrdering(t *testing.T, subscriptionName string) (message.Publisher, message.Subscriber) {
+	return newPubSub(
+		t,
+		true,
+		googlecloud.NewOrderingMarshaler(func(topic string, msg *message.Message) (string, error) {
+			return "ordering_key", nil
+		}),
+		googlecloud.NewOrderingUnmarshaler(func(orderingKey string, msg *message.Message) error {
+			return nil
+		}),
+		googlecloud.TopicSubscriptionNameWithSuffix(subscriptionName),
+	)
 }
 
 func TestPublishSubscribe(t *testing.T) {
@@ -69,6 +101,26 @@ func TestPublishSubscribe(t *testing.T) {
 	)
 }
 
+func TestPublishSubscribeOrdering(t *testing.T) {
+	t.Skip("skipping because the emulator does not currently redeliver nacked messages when ordering is enabled")
+
+	if testing.Short() {
+		t.Skip("skipping long tests")
+	}
+
+	tests.TestPubSub(
+		t,
+		tests.Features{
+			ConsumerGroups:      true,
+			ExactlyOnceDelivery: false,
+			GuaranteedOrder:     true,
+			Persistent:          true,
+		},
+		createPubSubWithOrdering,
+		createPubSubWithSubscriptionNameWithOrdering,
+	)
+}
+
 func TestSubscriberUnexpectedTopicForSubscription(t *testing.T) {
 	rand.Seed(time.Now().Unix())
 	testNumber := rand.Int()