SelectedHost: added Token function

This allows tokenAwareHostPolicy return token from ExecutableQuery.

Co-authored-by: Henrik Johansson <[email protected]>
Co-authored-by: Michał Matczuk <[email protected]>

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}
 	}
 }