From fb5d446fa3f7eba4aff7ace8f36cf425b8ea34fe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matczuk?= <michal@scylladb.com>
Date: Tue, 9 Oct 2018 21:47:47 +0200
Subject: [PATCH] SelectedHost: added Token function
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This allows tokenAwareHostPolicy return token from ExecutableQuery.

Co-authored-by: Henrik Johansson <henrik@scylladb.com>
Co-authored-by: Michał Matczuk <michal@scylladb.com>
---
 policies.go | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/policies.go b/policies.go
index 16beef280..f562676c8 100644
--- a/policies.go
+++ b/policies.go
@@ -312,16 +312,24 @@ type HostSelectionPolicy interface {
 // selection policy.
 type SelectedHost interface {
 	Info() *HostInfo
+	Token() token
 	Mark(error)
 }
 
-type selectedHost HostInfo
+type selectedHost struct {
+	info  *HostInfo
+	token token
+}
+
+func (host selectedHost) Info() *HostInfo {
+	return host.info
+}
 
-func (host *selectedHost) Info() *HostInfo {
-	return (*HostInfo)(host)
+func (host selectedHost) Token() token {
+	return host.token
 }
 
-func (host *selectedHost) Mark(err error) {}
+func (host selectedHost) Mark(err error) {}
 
 // NextHost is an iteration function over picked hosts
 type NextHost func() SelectedHost
@@ -361,7 +369,7 @@ func (r *roundRobinHostPolicy) Pick(qry ExecutableQuery) NextHost {
 		}
 		host := hosts[(pos)%uint32(len(hosts))]
 		i++
-		return (*selectedHost)(host)
+		return selectedHost{info: host}
 	}
 }
 
@@ -568,7 +576,7 @@ func (t *tokenAwareHostPolicy) Pick(qry ExecutableQuery) NextHost {
 
 			if h.IsUp() && t.fallback.IsLocal(h) {
 				used[h] = true
-				return (*selectedHost)(h)
+				return selectedHost{info: h, token: token}
 			}
 		}
 
@@ -715,6 +723,10 @@ func (host selectedHostPoolHost) Info() *HostInfo {
 	return host.info
 }
 
+func (host selectedHostPoolHost) Token() token {
+	return nil
+}
+
 func (host selectedHostPoolHost) Mark(err error) {
 	ip := host.info.ConnectAddress().String()
 
@@ -794,7 +806,7 @@ func (d *dcAwareRR) Pick(q ExecutableQuery) NextHost {
 		}
 		host := hosts[(pos)%uint32(len(hosts))]
 		i++
-		return (*selectedHost)(host)
+		return selectedHost{info: host}
 	}
 }