messages: WIP

End to end test.

Author: sougou

go/vt/tabletserver/message_manager.go

@@ -123,7 +123,7 @@ func NewMessageManager(tsv *TabletServer, table *TableInfo, connpool *ConnPool)
 		"select time_next, epoch, id, message from %v where time_next < %a order by time_next desc limit %a",
 		mm.name, ":time_next", ":max")
 	mm.ackQuery = buildParsedQuery(
-		"update %v set time_acked = %a, time_next = null where id in %a",
+		"update %v set time_acked = %a, time_next = null where id in %a and time_acked is null",
 		mm.name, ":time_acked", "::ids")
 	mm.postponeQuery = buildParsedQuery(
 		"update %v set time_next = %a+(%a<<epoch), epoch = epoch+1 where id in %a and time_acked is null",

go/vt/tabletserver/message_manager_test.go

@@ -538,7 +538,7 @@ func TestMMGenerate(t *testing.T) {
 	mm.Open()
 	defer mm.Close()
 	query, bv := mm.GenerateAckQuery([]string{"1", "2"})
-	wantQuery := "update foo set time_acked = :time_acked, time_next = null where id in ::ids"
+	wantQuery := "update foo set time_acked = :time_acked, time_next = null where id in ::ids and time_acked is null"
 	if query != wantQuery {
 		t.Errorf("GenerateAckQuery query: %s, want %s", query, wantQuery)
 	}

go/vt/tabletserver/tabletserver_test.go

@@ -1471,13 +1471,13 @@ func TestMessageAck(t *testing.T) {
 	}
 
 	_, err = tsv.MessageAck(ctx, &target, "msg", ids)
-	want = "error: query: select time_scheduled, id from msg where id in ('1', '2') limit 10001 for update is not supported"
+	want = "error: query: select time_scheduled, id from msg where id in ('1', '2') and time_acked is null limit 10001 for update is not supported"
 	if err == nil || err.Error() != want {
 		t.Errorf("tsv.MessageAck(invalid): %v, want %s", err, want)
 	}
 
 	db.AddQuery(
-		"select time_scheduled, id from msg where id in ('1', '2') limit 10001 for update",
+		"select time_scheduled, id from msg where id in ('1', '2') and time_acked is null limit 10001 for update",
 		&sqltypes.Result{
 			RowsAffected: 1,
 			Rows: [][]sqltypes.Value{{

go/vt/vtgate/scatter_conn.go

@@ -436,10 +436,16 @@ func (stc *ScatterConn) StreamExecuteMulti(
 func (stc *ScatterConn) MessageStream(ctx context.Context, keyspace string, shards []string, name string, sendReply func(*sqltypes.Result) error) error {
 	// mu is used to merge multiple sendReply calls into one.
 	var mu sync.Mutex
+	fieldSent := false
 	allErrors := stc.multiGo(ctx, "MessageStream", keyspace, shards, topodatapb.TabletType_MASTER, func(target *querypb.Target) error {
 		return stc.gateway.MessageStream(ctx, target, name, func(qr *sqltypes.Result) error {
 			mu.Lock()
 			defer mu.Unlock()
+			if fieldSent && len(qr.Rows) == 0 {
+				return nil
+			}
+			// First result is always the field info.
+			fieldSent = true
 			return sendReply(qr)
 		})
 	})

go/vt/vtgate/topo_utils.go

@@ -5,7 +5,6 @@
 package vtgate
 
 import (
-	"bytes"
 	"encoding/hex"
 	"fmt"
 
@@ -168,14 +167,14 @@ func mapExactShards(ctx context.Context, topoServ topo.SrvTopoServer, cell, keys
 	}
 	shardnum := 0
 	for shardnum < len(allShards) {
-		if bytes.Compare(kr.Start, []byte(allShards[shardnum].KeyRange.Start)) == 0 {
+		if key.KeyRangeStartEqual(kr, allShards[shardnum].KeyRange) {
 			break
 		}
 		shardnum++
 	}
 	for shardnum < len(allShards) {
 		shards = append(shards, allShards[shardnum].Name)
-		if bytes.Compare(kr.End, []byte(allShards[shardnum].KeyRange.End)) == 0 {
+		if key.KeyRangeEndEqual(kr, allShards[shardnum].KeyRange) {
 			return keyspace, shards, nil
 		}
 		shardnum++

go/vt/vtgate/vtgate.go

@@ -127,6 +127,7 @@ type VTGate struct {
 	logStreamExecuteKeyRanges   *logutil.ThrottledLogger
 	logStreamExecuteShards      *logutil.ThrottledLogger
 	logUpdateStream             *logutil.ThrottledLogger
+	logMessageStream            *logutil.ThrottledLogger
 }
 
 // RegisterVTGate defines the type of registration mechanism.
@@ -177,6 +178,7 @@ func Init(ctx context.Context, hc discovery.HealthCheck, topoServer topo.Server,
 		logStreamExecuteKeyRanges:   logutil.NewThrottledLogger("StreamExecuteKeyRanges", 5*time.Second),
 		logStreamExecuteShards:      logutil.NewThrottledLogger("StreamExecuteShards", 5*time.Second),
 		logUpdateStream:             logutil.NewThrottledLogger("UpdateStream", 5*time.Second),
+		logMessageStream:            logutil.NewThrottledLogger("MessageStream", 5*time.Second),
 	}
 
 	normalErrors = stats.NewMultiCounters("VtgateApiErrorCounts", []string{"Operation", "Keyspace", "DbType"})
@@ -828,7 +830,7 @@ func (vtg *VTGate) MessageStream(ctx context.Context, keyspace string, shard str
 			"TabletType":  ltt,
 			"MessageName": name,
 		}
-		logError(err, query, vtg.logUpdateStream)
+		logError(err, query, vtg.logMessageStream)
 	}
 	return formatError(err)
 }

py/vtdb/grpc_vtgate_client.py

@@ -249,6 +249,57 @@ def row_generator():
 
     return row_generator()
 
+  @vtgate_utils.exponential_backoff_retry((dbexceptions.ThrottledError,
+                                           dbexceptions.TransientError))
+  def message_stream(
+      self, keyspace, name,
+      shard=None, key_range=None,
+      effective_caller_id=None,
+      **kwargs):
+
+    try:
+      request = self.message_stream_request(
+          keyspace, shard, key_range,
+          name, effective_caller_id)
+      it = self.stub.MessageStream(request, self.timeout)
+      first_response = it.next()
+    except (grpc.RpcError, vtgate_utils.VitessError) as e:
+      raise _convert_exception(
+          e, 'MessageStream', name=name,
+          keyspace=keyspace)
+
+    fields, convs = self.build_conversions(first_response.result.fields)
+
+    def row_generator():
+      try:
+        for response in it:
+          for row in response.result.rows:
+            yield tuple(proto3_encoding.make_row(row, convs))
+      except Exception:
+        logging.exception('gRPC low-level error')
+        raise
+
+    return row_generator(), fields
+
+  @vtgate_utils.exponential_backoff_retry((dbexceptions.ThrottledError,
+                                           dbexceptions.TransientError))
+  def message_ack(
+      self,
+      name, ids,
+      keyspace=None, effective_caller_id=None,
+      **kwargs):
+
+    try:
+      request = self.message_ack_request(
+          keyspace, name, ids, effective_caller_id)
+      response = self.stub.MessageAck(request, self.timeout)
+    except (grpc.RpcError, vtgate_utils.VitessError) as e:
+      raise _convert_exception(
+          e, 'MessageAck', name=name, ids=ids,
+          keyspace=keyspace)
+
+    return response.count
+
 
 def _convert_exception(exc, *args, **kwargs):
   """This parses the protocol exceptions to the api interface exceptions.

py/vtdb/proto3_encoding.py

@@ -84,6 +84,13 @@ def make_row(row, convs):
   return converted_row
 
 
+def build_value(v):
+  """Build a proto value from any valid input."""
+  val = query_pb2.Value()
+  convert_value(v, val)
+  return val
+
+
 def convert_value(value, proto_value, allow_lists=False):
   """Convert a variable from python type to proto type+value.
 
@@ -547,6 +554,58 @@ def update_stream_request(self,
     self._add_caller_id(request, effective_caller_id)
     return request
 
+  def message_stream_request(self,
+                             keyspace_name,
+                             shard,
+                             key_range,
+                             name,
+                             effective_caller_id):
+    """Builds the right vtgate_pb2 MessageStreamRequest.
+
+    Args:
+      keyspace_name: keyspace to apply the query to.
+      shard: shard to ask for.
+      key_range: keyrange.KeyRange object.
+      name: message table name.
+      effective_caller_id: optional vtgate_client.CallerID.
+
+    Returns:
+      A vtgate_pb2.MessageStreamRequest object.
+    """
+    request = vtgate_pb2.MessageStreamRequest(keyspace=keyspace_name,
+                                              name=name,
+                                              shard=shard)
+    if key_range:
+      request.key_range.start = key_range.Start
+      request.key_range.end = key_range.End
+    self._add_caller_id(request, effective_caller_id)
+    return request
+
+  def message_ack_request(self,
+                          keyspace_name,
+                          name,
+                          ids,
+                          effective_caller_id):
+    """Builds the right vtgate_pb2 MessageAckRequest.
+
+    Args:
+      keyspace_name: keyspace to apply the query to.
+      name: message table name.
+      ids: list of message ids.
+      effective_caller_id: optional vtgate_client.CallerID.
+
+    Returns:
+      A vtgate_pb2.MessageAckRequest object.
+    """
+    vals = []
+    for v in ids:
+      vals.append(build_value(v))
+    request = vtgate_pb2.MessageAckRequest(keyspace=keyspace_name,
+                                           name=name,
+                                           ids=vals)
+    self._add_caller_id(request, effective_caller_id)
+    return request
+
   def stream_execute_request_and_name(self, sql, bind_variables, tablet_type,
                                       keyspace_name,
                                       shards,

py/vtdb/vtgate_client.py

@@ -389,3 +389,58 @@ def update_stream(self,
       dbexceptions.FatalError: this query should not be retried.
     """
     raise NotImplementedError('Child class needs to implement this')
+
+  def message_stream(self,
+                     keyspace, name,
+                     shard=None, key_range=None,
+                     effective_caller_id=None,
+                     **kwargs):
+    """Asks for a message stream.
+
+    Args:
+      keyspace: the keyspace of the message table.
+      name: the name of the message table.
+      shard: the shard name to listen for.
+        Incompatible with key_range.
+      key_range: the key range to listen for.
+        Incompatible with shard.
+      effective_caller_id: CallerID object.
+      **kwargs: implementation specific parameters.
+
+    Returns:
+      A (row generator, fields) pair.
+
+    Raises:
+      dbexceptions.TimeoutError: for connection timeout.
+      dbexceptions.TransientError: the server is overloaded, and this query
+        is asked to back off.
+      dbexceptions.DatabaseError: generic database error.
+      dbexceptions.FatalError: this query should not be retried.
+    """
+    raise NotImplementedError('Child class needs to implement this')
+
+  def message_ack(self,
+                  name, ids,
+                  keyspace=None, effective_caller_id=None,
+                  **kwargs):
+    """Acks a list of messages.
+
+    Args:
+      name: the name of the message table.
+      ids: list of message ids to ack.
+      keyspace: the keyspace of the message table.
+        Not required if table can be auto-resolved.
+      effective_caller_id: CallerID object.
+      **kwargs: implementation specific parameters.
+
+    Returns:
+      The number of rows acked.
+
+    Raises:
+      dbexceptions.TimeoutError: for connection timeout.
+      dbexceptions.TransientError: the server is overloaded, and this query
+        is asked to back off.
+      dbexceptions.DatabaseError: generic database error.
+      dbexceptions.FatalError: this query should not be retried.
+    """
+    raise NotImplementedError('Child class needs to implement this')

test/config.json

@@ -467,6 +467,15 @@
 			"RetryMax": 0,
 			"Tags": []
 		},
+		"messaging": {
+			"File": "messaging.py",
+			"Args": [],
+			"Command": [],
+			"Manual": false,
+			"Shard": 4,
+			"RetryMax": 0,
+			"Tags": []
+		},
 		"vttest_sample": {
 			"File": "vttest_sample_test.py",
 			"Args": [],