Add Transactionalmessage interface for request routing (segmentio#744)

* Add TransactionalMessage interface. Add TransactionMessage interface so the Transport automatically routes messages to Transaction Coordinators if required.
jiujuan · Sep 20, 2021 · 392c450 · 392c450
1 parent bedcf4a
commit 392c450
Show file tree

Hide file tree

Showing 7 changed files with 38 additions and 0 deletions.
diff --git a/protocol/addoffsetstotxn/addoffsetstotxn.go b/protocol/addoffsetstotxn/addoffsetstotxn.go
@@ -19,6 +19,10 @@ type Request struct {
 
 func (r *Request) ApiKey() protocol.ApiKey { return protocol.AddOffsetsToTxn }
 
+func (r *Request) Transaction() string { return r.TransactionalID }
+
+var _ protocol.TransactionalMessage = (*Request)(nil)
+
 type Response struct {
 	// We need at least one tagged field to indicate that this is a "flexible" message
 	// type.

diff --git a/protocol/addpartitionstotxn/addpartitionstotxn.go b/protocol/addpartitionstotxn/addpartitionstotxn.go
@@ -28,6 +28,10 @@ type RequestTopic struct {
 
 func (r *Request) ApiKey() protocol.ApiKey { return protocol.AddPartitionsToTxn }
 
+func (r *Request) Transaction() string { return r.TransactionalID }
+
+var _ protocol.TransactionalMessage = (*Request)(nil)
+
 type Response struct {
 	// We need at least one tagged field to indicate that this is a "flexible" message
 	// type.

diff --git a/protocol/endtxn/endtxn.go b/protocol/endtxn/endtxn.go
@@ -19,6 +19,10 @@ type Request struct {
 
 func (r *Request) ApiKey() protocol.ApiKey { return protocol.EndTxn }
 
+func (r *Request) Transaction() string { return r.TransactionalID }
+
+var _ protocol.TransactionalMessage = (*Request)(nil)
+
 type Response struct {
 	// We need at least one tagged field to indicate that this is a "flexible" message
 	// type.

diff --git a/protocol/initproducerid/initproducerid.go b/protocol/initproducerid/initproducerid.go
@@ -19,6 +19,10 @@ type Request struct {
 
 func (r *Request) ApiKey() protocol.ApiKey { return protocol.InitProducerId }
 
+func (r *Request) Transaction() string { return r.TransactionalID }
+
+var _ protocol.TransactionalMessage = (*Request)(nil)
+
 type Response struct {
 	// We need at least one tagged field to indicate that this is a "flexible" message
 	// type.

diff --git a/protocol/protocol.go b/protocol/protocol.go
@@ -438,6 +438,14 @@ type GroupMessage interface {
 	Group() string
 }
 
+// TransactionalMessage is an extension of the Message interface implemented by some
+// request types to inform the program that they should be reouted to a transaction
+// coordinator.
+type TransactionalMessage interface {
+	// Returns the transactional id configured on the message.
+	Transaction() string
+}
+
 // PreparedMessage is an extension of the Message interface implemented by some
 // request types which may need to run some pre-processing on their state before
 // being sent.

diff --git a/protocol/txnoffsetcommit/txnoffsetcommit.go b/protocol/txnoffsetcommit/txnoffsetcommit.go
@@ -43,6 +43,10 @@ type RequestPartition struct {
 
 func (r *Request) ApiKey() protocol.ApiKey { return protocol.TxnOffsetCommit }
 
+func (r *Request) Group() string { return r.GroupID }
+
+var _ protocol.GroupMessage = (*Request)(nil)
+
 type Response struct {
 	// We need at least one tagged field to indicate that this is a "flexible" message
 	// type.

diff --git a/transport.go b/transport.go
@@ -649,6 +649,16 @@ func (p *connPool) sendRequest(ctx context.Context, req Request, state connPoolS
 			return reject(err)
 		}
 		brokerID = r.(*findcoordinator.Response).NodeID
+	case protocol.TransactionalMessage:
+		p := p.sendRequest(ctx, &findcoordinator.Request{
+			Key:     m.Transaction(),
+			KeyType: int8(CoordinatorKeyTypeTransaction),
+		}, state)
+		r, err := p.await(ctx)
+		if err != nil {
+			return reject(err)
+		}
+		brokerID = r.(*findcoordinator.Response).NodeID
 	}
 
 	var c *conn