Use tablets in TokenAwareHostPolicy

Add mechanism to parse system.tablets periodically. In TokenAwareHostPolicy
check if keyspace uses tablets if so try to use them to find replicas.

Make shard awareness work when using tablets.

Author: sylwiaszunejko

conn.go

@@ -208,7 +208,8 @@ type Conn struct {
 
 	timeouts int64
 
-	logger StdLogger
+	logger           StdLogger
+	tabletsRoutingV1 bool
 }
 
 // connect establishes a connection to a Cassandra node using session's connection config.
@@ -724,6 +725,9 @@ func (c *Conn) recv(ctx context.Context) error {
 	} else if head.stream == -1 {
 		// TODO: handle cassandra event frames, we shouldnt get any currently
 		framer := newFramerWithExts(c.compressor, c.version, c.cqlProtoExts)
+		c.mu.Lock()
+		c.tabletsRoutingV1 = framer.tabletsRoutingV1
+		c.mu.Unlock()
 		if err := framer.readFrame(c, &head); err != nil {
 			return err
 		}
@@ -733,6 +737,9 @@ func (c *Conn) recv(ctx context.Context) error {
 		// reserved stream that we dont use, probably due to a protocol error
 		// or a bug in Cassandra, this should be an error, parse it and return.
 		framer := newFramerWithExts(c.compressor, c.version, c.cqlProtoExts)
+		c.mu.Lock()
+		c.tabletsRoutingV1 = framer.tabletsRoutingV1
+		c.mu.Unlock()
 		if err := framer.readFrame(c, &head); err != nil {
 			return err
 		}
@@ -1069,6 +1076,9 @@ func (c *Conn) exec(ctx context.Context, req frameBuilder, tracer Tracer) (*fram
 
 	// resp is basically a waiting semaphore protecting the framer
 	framer := newFramerWithExts(c.compressor, c.version, c.cqlProtoExts)
+	c.mu.Lock()
+	c.tabletsRoutingV1 = framer.tabletsRoutingV1
+	c.mu.Unlock()
 
 	call := &callReq{
 		timeout:  make(chan struct{}),
@@ -1453,6 +1463,63 @@ func (c *Conn) executeQuery(ctx context.Context, qry *Query) *Iter {
 		return &Iter{err: err}
 	}
 
+	if len(framer.customPayload) > 0 {
+		if tabletInfo, ok := framer.customPayload["tablets-routing-v1"]; ok {
+			var firstToken string
+			var lastToken string
+			var replicas [][]interface{}
+			tabletInfoValue := []interface{}{&firstToken, &lastToken, &replicas}
+			Unmarshal(TupleTypeInfo{
+				NativeType: NativeType{proto: c.version, typ: TypeTuple},
+				Elems: []TypeInfo{
+					NativeType{typ: TypeBigInt},
+					NativeType{typ: TypeBigInt},
+					CollectionType{
+						NativeType: NativeType{proto: c.version, typ: TypeList},
+						Elem: TupleTypeInfo{
+							NativeType: NativeType{proto: c.version, typ: TypeTuple},
+							Elems: []TypeInfo{
+								NativeType{proto: c.version, typ: TypeUUID},
+								NativeType{proto: c.version, typ: TypeInt},
+							}},
+					},
+				},
+			}, tabletInfo, tabletInfoValue)
+
+			tablet := TabletInfo{}
+			tablet.firstToken, err = strconv.ParseInt(firstToken, 10, 64)
+			if err != nil {
+				return &Iter{err: err}
+			}
+			tablet.lastToken, err = strconv.ParseInt(lastToken, 10, 64)
+			if err != nil {
+				return &Iter{err: err}
+			}
+
+			tabletReplicas := make([]ReplicaInfo, 0, len(replicas))
+			for _, replica := range replicas {
+				if len(replica) != 2 {
+					return &Iter{err: err}
+				}
+				if hostId, ok := replica[0].(UUID); ok {
+					if shardId, ok := replica[1].(int); ok {
+						repInfo := ReplicaInfo{hostId, shardId}
+						tabletReplicas = append(tabletReplicas, repInfo)
+					} else {
+						return &Iter{err: err}
+					}
+				} else {
+					return &Iter{err: err}
+				}
+			}
+			tablet.replicas = tabletReplicas
+			tablet.keyspaceName = qry.routingInfo.keyspace
+			tablet.tableName = qry.routingInfo.table
+
+			addTablet(c.session.hostSource, &tablet)
+		}
+	}
+
 	if len(framer.traceID) > 0 && qry.trace != nil {
 		qry.trace.Trace(framer.traceID)
 	}

connectionpool.go

@@ -26,6 +26,12 @@ type SetPartitioner interface {
 	SetPartitioner(partitioner string)
 }
 
+// interface to implement to receive the tablets value
+// Experimental, this interface and use may change
+type SetTablets interface {
+	SetTablets(tablets []*TabletInfo)
+}
+
 func setupTLSConfig(sslOpts *SslOptions) (*tls.Config, error) {
 	//  Config.InsecureSkipVerify | EnableHostVerification | Result
 	//  Config is nil             | true                   | verify host
@@ -312,7 +318,7 @@ func newHostConnPool(session *Session, host *HostInfo, port, size int,
 }
 
 // Pick a connection from this connection pool for the given query.
-func (pool *hostConnPool) Pick(token token) *Conn {
+func (pool *hostConnPool) Pick(token token, keyspace string, table string) *Conn {
 	pool.mu.RLock()
 	defer pool.mu.RUnlock()
 
@@ -330,7 +336,7 @@ func (pool *hostConnPool) Pick(token token) *Conn {
 		}
 	}
 
-	return pool.connPicker.Pick(token)
+	return pool.connPicker.Pick(token, keyspace, table)
 }
 
 // Size returns the number of connections currently active in the pool

connpicker.go

@@ -7,7 +7,7 @@ import (
 )
 
 type ConnPicker interface {
-	Pick(token) *Conn
+	Pick(token, string, string) *Conn
 	Put(*Conn)
 	Remove(conn *Conn)
 	Size() (int, int)
@@ -65,7 +65,7 @@ func (p *defaultConnPicker) Size() (int, int) {
 	return size, p.size - size
 }
 
-func (p *defaultConnPicker) Pick(token) *Conn {
+func (p *defaultConnPicker) Pick(token, string, string) *Conn {
 	pos := int(atomic.AddUint32(&p.pos, 1) - 1)
 	size := len(p.conns)
 
@@ -104,7 +104,7 @@ func (*defaultConnPicker) NextShard() (shardID, nrShards int) {
 // to the point where we have first connection.
 type nopConnPicker struct{}
 
-func (nopConnPicker) Pick(token) *Conn {
+func (nopConnPicker) Pick(token, string, string) *Conn {
 	return nil
 }
 

frame.go

@@ -367,6 +367,7 @@ type framer struct {
 
 	flagLWT               int
 	rateLimitingErrorCode int
+	tabletsRoutingV1      bool
 }
 
 func newFramer(compressor Compressor, version byte) *framer {
@@ -398,6 +399,8 @@ func newFramer(compressor Compressor, version byte) *framer {
 	f.header = nil
 	f.traceID = nil
 
+	f.tabletsRoutingV1 = false
+
 	return f
 }
 
@@ -435,6 +438,7 @@ func newFramerWithExts(compressor Compressor, version byte, cqlProtoExts []cqlPr
 					tabletsRoutingV1, tabletsRoutingV1Ext{}))
 			return f
 		}
+		f.tabletsRoutingV1 = true
 	}
 
 	return f

host_source.go

@@ -472,12 +472,151 @@ func (h *HostInfo) ScyllaShardAwarePortTLS() uint16 {
 	return h.scyllaShardAwarePortTLS
 }
 
+// Experimental, this interface and use may change
+type ReplicaInfo struct {
+	hostId  UUID
+	shardId int
+}
+
+// Experimental, this interface and use may change
+type TabletInfo struct {
+	mu           sync.RWMutex
+	keyspaceName string
+	tableName    string
+	firstToken   int64
+	lastToken    int64
+	replicas     []ReplicaInfo
+}
+
+func (t *TabletInfo) KeyspaceName() string {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	return t.keyspaceName
+}
+
+func (t *TabletInfo) FirstToken() int64 {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	return t.firstToken
+}
+
+func (t *TabletInfo) LastToken() int64 {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	return t.lastToken
+}
+
+func (t *TabletInfo) TableName() string {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	return t.tableName
+}
+
+func (t *TabletInfo) Replicas() []ReplicaInfo {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	return t.replicas
+}
+
+// Search for place in tablets table with specific Keyspace and Table name
+func findTablets(tablets []*TabletInfo, k string, t string) (int, int) {
+	l := -1
+	r := -1
+	for i, tablet := range tablets {
+		if tablet.KeyspaceName() == k && tablet.TableName() == t {
+			if l == -1 {
+				l = i
+			}
+			r = i
+		} else if l != -1 {
+			break
+		}
+	}
+
+	return l, r
+}
+
+func addTabletToTabletsList(tablets []*TabletInfo, tablet *TabletInfo) []*TabletInfo {
+	l, r := findTablets(tablets, tablet.keyspaceName, tablet.tableName)
+	if l == -1 && r == -1 {
+		l = 0
+		r = 0
+	} else {
+		r = r + 1
+	}
+
+	l1, r1 := l, r
+	l2, r2 := l1, r1
+
+	// find first overlaping range
+	for l1 < r1 {
+		mid := (l1 + r1) / 2
+		if tablets[mid].FirstToken() < tablet.FirstToken() {
+			l1 = mid + 1
+		} else {
+			r1 = mid
+		}
+	}
+	start := l1
+
+	if start > l && tablets[start-1].LastToken() > tablet.FirstToken() {
+		start = start - 1
+	}
+
+	// find last overlaping range
+	for l2 < r2 {
+		mid := (l2 + r2) / 2
+		if tablets[mid].LastToken() < tablet.LastToken() {
+			l2 = mid + 1
+		} else {
+			r2 = mid
+		}
+	}
+	end := l2
+	if end < r && tablets[end].FirstToken() >= tablet.LastToken() {
+		end = end - 1
+	}
+	if end == len(tablets) {
+		end = end - 1
+	}
+
+	updated_tablets := tablets
+	if start <= end {
+		// Delete elements from index start to end
+		updated_tablets = append(tablets[:start], tablets[end+1:]...)
+	}
+	// Insert tablet element at index start
+	updated_tablets2 := append(updated_tablets[:start], append([]*TabletInfo{tablet}, updated_tablets[start:]...)...)
+	return updated_tablets2
+}
+
+// Search for place in tablets table for token starting from index l to index r
+func findTabletForToken(tablets []*TabletInfo, token token, l int, r int) *TabletInfo {
+	for l < r {
+		var m int
+		if r*l > 0 {
+			m = l + (r-l)/2
+		} else {
+			m = (r + l) / 2
+		}
+		if int64Token(tablets[m].LastToken()).Less(token) {
+			l = m + 1
+		} else {
+			r = m
+		}
+	}
+
+	return tablets[l]
+}
+
 // Polls system.peers at a specific interval to find new hosts
 type ringDescriber struct {
 	session         *Session
 	mu              sync.Mutex
 	prevHosts       []*HostInfo
 	prevPartitioner string
+	// Experimental, this interface and use may change
+	prevTablets []*TabletInfo
 }
 
 // Returns true if we are using system_schema.keyspaces instead of system.schema_keyspaces
@@ -835,6 +974,21 @@ func refreshRing(r *ringDescriber) error {
 
 	r.session.metadata.setPartitioner(partitioner)
 	r.session.policy.SetPartitioner(partitioner)
+
+	return nil
+}
+
+// Experimental, this interface and use may change
+func addTablet(r *ringDescriber, tablet *TabletInfo) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	tablets := r.session.getTablets()
+	tablets = addTabletToTabletsList(tablets, tablet)
+
+	r.session.ring.setTablets(tablets)
+	r.session.policy.SetTablets(tablets)
+
 	return nil
 }