policies.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886
  1. // Copyright (c) 2012 The gocql Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. //This file will be the future home for more policies
  5. package gocql
  6. import (
  7. "context"
  8. "errors"
  9. "fmt"
  10. "math"
  11. "math/rand"
  12. "net"
  13. "sync"
  14. "sync/atomic"
  15. "time"
  16. "github.com/hailocab/go-hostpool"
  17. )
  18. // cowHostList implements a copy on write host list, its equivalent type is []*HostInfo
  19. type cowHostList struct {
  20. list atomic.Value
  21. mu sync.Mutex
  22. }
  23. func (c *cowHostList) String() string {
  24. return fmt.Sprintf("%+v", c.get())
  25. }
  26. func (c *cowHostList) get() []*HostInfo {
  27. // TODO(zariel): should we replace this with []*HostInfo?
  28. l, ok := c.list.Load().(*[]*HostInfo)
  29. if !ok {
  30. return nil
  31. }
  32. return *l
  33. }
  34. func (c *cowHostList) set(list []*HostInfo) {
  35. c.mu.Lock()
  36. c.list.Store(&list)
  37. c.mu.Unlock()
  38. }
  39. // add will add a host if it not already in the list
  40. func (c *cowHostList) add(host *HostInfo) bool {
  41. c.mu.Lock()
  42. l := c.get()
  43. if n := len(l); n == 0 {
  44. l = []*HostInfo{host}
  45. } else {
  46. newL := make([]*HostInfo, n+1)
  47. for i := 0; i < n; i++ {
  48. if host.Equal(l[i]) {
  49. c.mu.Unlock()
  50. return false
  51. }
  52. newL[i] = l[i]
  53. }
  54. newL[n] = host
  55. l = newL
  56. }
  57. c.list.Store(&l)
  58. c.mu.Unlock()
  59. return true
  60. }
  61. func (c *cowHostList) update(host *HostInfo) {
  62. c.mu.Lock()
  63. l := c.get()
  64. if len(l) == 0 {
  65. c.mu.Unlock()
  66. return
  67. }
  68. found := false
  69. newL := make([]*HostInfo, len(l))
  70. for i := range l {
  71. if host.Equal(l[i]) {
  72. newL[i] = host
  73. found = true
  74. } else {
  75. newL[i] = l[i]
  76. }
  77. }
  78. if found {
  79. c.list.Store(&newL)
  80. }
  81. c.mu.Unlock()
  82. }
  83. func (c *cowHostList) remove(ip net.IP) bool {
  84. c.mu.Lock()
  85. l := c.get()
  86. size := len(l)
  87. if size == 0 {
  88. c.mu.Unlock()
  89. return false
  90. }
  91. found := false
  92. newL := make([]*HostInfo, 0, size)
  93. for i := 0; i < len(l); i++ {
  94. if !l[i].ConnectAddress().Equal(ip) {
  95. newL = append(newL, l[i])
  96. } else {
  97. found = true
  98. }
  99. }
  100. if !found {
  101. c.mu.Unlock()
  102. return false
  103. }
  104. newL = newL[: size-1 : size-1]
  105. c.list.Store(&newL)
  106. c.mu.Unlock()
  107. return true
  108. }
  109. // RetryableQuery is an interface that represents a query or batch statement that
  110. // exposes the correct functions for the retry policy logic to evaluate correctly.
  111. type RetryableQuery interface {
  112. Attempts() int
  113. SetConsistency(c Consistency)
  114. GetConsistency() Consistency
  115. Context() context.Context
  116. }
  117. type RetryType uint16
  118. const (
  119. Retry RetryType = 0x00 // retry on same connection
  120. RetryNextHost RetryType = 0x01 // retry on another connection
  121. Ignore RetryType = 0x02 // ignore error and return result
  122. Rethrow RetryType = 0x03 // raise error and stop retrying
  123. )
  124. // ErrUnknownRetryType is returned if the retry policy returns a retry type
  125. // unknown to the query executor.
  126. var ErrUnknownRetryType = errors.New("unknown retry type returned by retry policy")
  127. // RetryPolicy interface is used by gocql to determine if a query can be attempted
  128. // again after a retryable error has been received. The interface allows gocql
  129. // users to implement their own logic to determine if a query can be attempted
  130. // again.
  131. //
  132. // See SimpleRetryPolicy as an example of implementing and using a RetryPolicy
  133. // interface.
  134. type RetryPolicy interface {
  135. Attempt(RetryableQuery) bool
  136. GetRetryType(error) RetryType
  137. }
  138. // SimpleRetryPolicy has simple logic for attempting a query a fixed number of times.
  139. //
  140. // See below for examples of usage:
  141. //
  142. // //Assign to the cluster
  143. // cluster.RetryPolicy = &gocql.SimpleRetryPolicy{NumRetries: 3}
  144. //
  145. // //Assign to a query
  146. // query.RetryPolicy(&gocql.SimpleRetryPolicy{NumRetries: 1})
  147. //
  148. type SimpleRetryPolicy struct {
  149. NumRetries int //Number of times to retry a query
  150. }
  151. // Attempt tells gocql to attempt the query again based on query.Attempts being less
  152. // than the NumRetries defined in the policy.
  153. func (s *SimpleRetryPolicy) Attempt(q RetryableQuery) bool {
  154. return q.Attempts() <= s.NumRetries
  155. }
  156. func (s *SimpleRetryPolicy) GetRetryType(err error) RetryType {
  157. return RetryNextHost
  158. }
  159. // ExponentialBackoffRetryPolicy sleeps between attempts
  160. type ExponentialBackoffRetryPolicy struct {
  161. NumRetries int
  162. Min, Max time.Duration
  163. }
  164. func (e *ExponentialBackoffRetryPolicy) Attempt(q RetryableQuery) bool {
  165. if q.Attempts() > e.NumRetries {
  166. return false
  167. }
  168. time.Sleep(e.napTime(q.Attempts()))
  169. return true
  170. }
  171. // used to calculate exponentially growing time
  172. func getExponentialTime(min time.Duration, max time.Duration, attempts int) time.Duration {
  173. if min <= 0 {
  174. min = 100 * time.Millisecond
  175. }
  176. if max <= 0 {
  177. max = 10 * time.Second
  178. }
  179. minFloat := float64(min)
  180. napDuration := minFloat * math.Pow(2, float64(attempts-1))
  181. // add some jitter
  182. napDuration += rand.Float64()*minFloat - (minFloat / 2)
  183. if napDuration > float64(max) {
  184. return time.Duration(max)
  185. }
  186. return time.Duration(napDuration)
  187. }
  188. func (e *ExponentialBackoffRetryPolicy) GetRetryType(err error) RetryType {
  189. return RetryNextHost
  190. }
  191. // DowngradingConsistencyRetryPolicy: Next retry will be with the next consistency level
  192. // provided in the slice
  193. //
  194. // On a read timeout: the operation is retried with the next provided consistency
  195. // level.
  196. //
  197. // On a write timeout: if the operation is an :attr:`~.UNLOGGED_BATCH`
  198. // and at least one replica acknowledged the write, the operation is
  199. // retried with the next consistency level. Furthermore, for other
  200. // write types, if at least one replica acknowledged the write, the
  201. // timeout is ignored.
  202. //
  203. // On an unavailable exception: if at least one replica is alive, the
  204. // operation is retried with the next provided consistency level.
  205. type DowngradingConsistencyRetryPolicy struct {
  206. ConsistencyLevelsToTry []Consistency
  207. }
  208. func (d *DowngradingConsistencyRetryPolicy) Attempt(q RetryableQuery) bool {
  209. currentAttempt := q.Attempts()
  210. if currentAttempt > len(d.ConsistencyLevelsToTry) {
  211. return false
  212. } else if currentAttempt > 0 {
  213. q.SetConsistency(d.ConsistencyLevelsToTry[currentAttempt-1])
  214. if gocqlDebug {
  215. Logger.Printf("%T: set consistency to %q\n",
  216. d,
  217. d.ConsistencyLevelsToTry[currentAttempt-1])
  218. }
  219. }
  220. return true
  221. }
  222. func (d *DowngradingConsistencyRetryPolicy) GetRetryType(err error) RetryType {
  223. switch t := err.(type) {
  224. case *RequestErrUnavailable:
  225. if t.Alive > 0 {
  226. return Retry
  227. }
  228. return Rethrow
  229. case *RequestErrWriteTimeout:
  230. if t.WriteType == "SIMPLE" || t.WriteType == "BATCH" || t.WriteType == "COUNTER" {
  231. if t.Received > 0 {
  232. return Ignore
  233. }
  234. return Rethrow
  235. }
  236. if t.WriteType == "UNLOGGED_BATCH" {
  237. return Retry
  238. }
  239. return Rethrow
  240. case *RequestErrReadTimeout:
  241. return Retry
  242. default:
  243. return RetryNextHost
  244. }
  245. }
  246. func (e *ExponentialBackoffRetryPolicy) napTime(attempts int) time.Duration {
  247. return getExponentialTime(e.Min, e.Max, attempts)
  248. }
  249. type HostStateNotifier interface {
  250. AddHost(host *HostInfo)
  251. RemoveHost(host *HostInfo)
  252. HostUp(host *HostInfo)
  253. HostDown(host *HostInfo)
  254. }
  255. type KeyspaceUpdateEvent struct {
  256. Keyspace string
  257. Change string
  258. }
  259. // HostSelectionPolicy is an interface for selecting
  260. // the most appropriate host to execute a given query.
  261. type HostSelectionPolicy interface {
  262. HostStateNotifier
  263. SetPartitioner
  264. KeyspaceChanged(KeyspaceUpdateEvent)
  265. Init(*Session)
  266. IsLocal(host *HostInfo) bool
  267. //Pick returns an iteration function over selected hosts
  268. Pick(ExecutableQuery) NextHost
  269. }
  270. // SelectedHost is an interface returned when picking a host from a host
  271. // selection policy.
  272. type SelectedHost interface {
  273. Info() *HostInfo
  274. Mark(error)
  275. }
  276. type selectedHost HostInfo
  277. func (host *selectedHost) Info() *HostInfo {
  278. return (*HostInfo)(host)
  279. }
  280. func (host *selectedHost) Mark(err error) {}
  281. // NextHost is an iteration function over picked hosts
  282. type NextHost func() SelectedHost
  283. // RoundRobinHostPolicy is a round-robin load balancing policy, where each host
  284. // is tried sequentially for each query.
  285. func RoundRobinHostPolicy() HostSelectionPolicy {
  286. return &roundRobinHostPolicy{}
  287. }
  288. type roundRobinHostPolicy struct {
  289. hosts cowHostList
  290. pos uint32
  291. mu sync.RWMutex
  292. }
  293. func (r *roundRobinHostPolicy) IsLocal(*HostInfo) bool { return true }
  294. func (r *roundRobinHostPolicy) KeyspaceChanged(KeyspaceUpdateEvent) {}
  295. func (r *roundRobinHostPolicy) SetPartitioner(partitioner string) {}
  296. func (r *roundRobinHostPolicy) Init(*Session) {}
  297. func (r *roundRobinHostPolicy) Pick(qry ExecutableQuery) NextHost {
  298. // i is used to limit the number of attempts to find a host
  299. // to the number of hosts known to this policy
  300. var i int
  301. return func() SelectedHost {
  302. hosts := r.hosts.get()
  303. if len(hosts) == 0 {
  304. return nil
  305. }
  306. // always increment pos to evenly distribute traffic in case of
  307. // failures
  308. pos := atomic.AddUint32(&r.pos, 1) - 1
  309. if i >= len(hosts) {
  310. return nil
  311. }
  312. host := hosts[(pos)%uint32(len(hosts))]
  313. i++
  314. return (*selectedHost)(host)
  315. }
  316. }
  317. func (r *roundRobinHostPolicy) AddHost(host *HostInfo) {
  318. r.hosts.add(host)
  319. }
  320. func (r *roundRobinHostPolicy) RemoveHost(host *HostInfo) {
  321. r.hosts.remove(host.ConnectAddress())
  322. }
  323. func (r *roundRobinHostPolicy) HostUp(host *HostInfo) {
  324. r.AddHost(host)
  325. }
  326. func (r *roundRobinHostPolicy) HostDown(host *HostInfo) {
  327. r.RemoveHost(host)
  328. }
  329. func ShuffleReplicas() func(*tokenAwareHostPolicy) {
  330. return func(t *tokenAwareHostPolicy) {
  331. t.shuffleReplicas = true
  332. }
  333. }
  334. // TokenAwareHostPolicy is a token aware host selection policy, where hosts are
  335. // selected based on the partition key, so queries are sent to the host which
  336. // owns the partition. Fallback is used when routing information is not available.
  337. func TokenAwareHostPolicy(fallback HostSelectionPolicy, opts ...func(*tokenAwareHostPolicy)) HostSelectionPolicy {
  338. p := &tokenAwareHostPolicy{fallback: fallback}
  339. for _, opt := range opts {
  340. opt(p)
  341. }
  342. return p
  343. }
  344. type keyspaceMeta struct {
  345. replicas map[string]map[token][]*HostInfo
  346. }
  347. type tokenAwareHostPolicy struct {
  348. hosts cowHostList
  349. mu sync.RWMutex
  350. partitioner string
  351. fallback HostSelectionPolicy
  352. session *Session
  353. tokenRing atomic.Value // *tokenRing
  354. keyspaces atomic.Value // *keyspaceMeta
  355. shuffleReplicas bool
  356. }
  357. func (t *tokenAwareHostPolicy) Init(s *Session) {
  358. t.session = s
  359. }
  360. func (t *tokenAwareHostPolicy) IsLocal(host *HostInfo) bool {
  361. return t.fallback.IsLocal(host)
  362. }
  363. func (t *tokenAwareHostPolicy) KeyspaceChanged(update KeyspaceUpdateEvent) {
  364. t.updateKeyspaceMetadata(update.Keyspace)
  365. }
  366. func (t *tokenAwareHostPolicy) updateKeyspaceMetadata(keyspace string) {
  367. meta, _ := t.keyspaces.Load().(*keyspaceMeta)
  368. var size = 1
  369. if meta != nil {
  370. size = len(meta.replicas)
  371. }
  372. newMeta := &keyspaceMeta{
  373. replicas: make(map[string]map[token][]*HostInfo, size),
  374. }
  375. ks, err := t.session.KeyspaceMetadata(keyspace)
  376. if err == nil {
  377. strat := getStrategy(ks)
  378. if strat != nil {
  379. tr, _ := t.tokenRing.Load().(*tokenRing)
  380. if tr != nil {
  381. newMeta.replicas[keyspace] = strat.replicaMap(t.hosts.get(), tr.tokens)
  382. }
  383. }
  384. }
  385. if meta != nil {
  386. for ks, replicas := range meta.replicas {
  387. if ks != keyspace {
  388. newMeta.replicas[ks] = replicas
  389. }
  390. }
  391. }
  392. t.keyspaces.Store(newMeta)
  393. }
  394. func (t *tokenAwareHostPolicy) SetPartitioner(partitioner string) {
  395. t.mu.Lock()
  396. defer t.mu.Unlock()
  397. if t.partitioner != partitioner {
  398. t.fallback.SetPartitioner(partitioner)
  399. t.partitioner = partitioner
  400. t.resetTokenRing(partitioner)
  401. }
  402. }
  403. func (t *tokenAwareHostPolicy) AddHost(host *HostInfo) {
  404. t.HostUp(host)
  405. if t.session != nil { // disable for unit tests
  406. t.updateKeyspaceMetadata(t.session.cfg.Keyspace)
  407. }
  408. }
  409. func (t *tokenAwareHostPolicy) RemoveHost(host *HostInfo) {
  410. t.HostDown(host)
  411. if t.session != nil { // disable for unit tests
  412. t.updateKeyspaceMetadata(t.session.cfg.Keyspace)
  413. }
  414. }
  415. func (t *tokenAwareHostPolicy) HostUp(host *HostInfo) {
  416. t.hosts.add(host)
  417. t.fallback.AddHost(host)
  418. t.mu.RLock()
  419. partitioner := t.partitioner
  420. t.mu.RUnlock()
  421. t.resetTokenRing(partitioner)
  422. }
  423. func (t *tokenAwareHostPolicy) HostDown(host *HostInfo) {
  424. t.hosts.remove(host.ConnectAddress())
  425. t.fallback.RemoveHost(host)
  426. t.mu.RLock()
  427. partitioner := t.partitioner
  428. t.mu.RUnlock()
  429. t.resetTokenRing(partitioner)
  430. }
  431. func (t *tokenAwareHostPolicy) resetTokenRing(partitioner string) {
  432. if partitioner == "" {
  433. // partitioner not yet set
  434. return
  435. }
  436. // create a new token ring
  437. hosts := t.hosts.get()
  438. tokenRing, err := newTokenRing(partitioner, hosts)
  439. if err != nil {
  440. Logger.Printf("Unable to update the token ring due to error: %s", err)
  441. return
  442. }
  443. // replace the token ring
  444. t.tokenRing.Store(tokenRing)
  445. }
  446. func (t *tokenAwareHostPolicy) getReplicas(keyspace string, token token) ([]*HostInfo, bool) {
  447. meta, _ := t.keyspaces.Load().(*keyspaceMeta)
  448. if meta == nil {
  449. return nil, false
  450. }
  451. replicas, ok := meta.replicas[keyspace][token]
  452. return replicas, ok
  453. }
  454. func (t *tokenAwareHostPolicy) Pick(qry ExecutableQuery) NextHost {
  455. if qry == nil {
  456. return t.fallback.Pick(qry)
  457. }
  458. routingKey, err := qry.GetRoutingKey()
  459. if err != nil {
  460. return t.fallback.Pick(qry)
  461. } else if routingKey == nil {
  462. return t.fallback.Pick(qry)
  463. }
  464. tr, _ := t.tokenRing.Load().(*tokenRing)
  465. if tr == nil {
  466. return t.fallback.Pick(qry)
  467. }
  468. primaryEndpoint, token := tr.GetHostForPartitionKey(routingKey)
  469. if primaryEndpoint == nil || token == nil {
  470. return t.fallback.Pick(qry)
  471. }
  472. replicas, ok := t.getReplicas(qry.Keyspace(), token)
  473. if !ok {
  474. replicas = []*HostInfo{primaryEndpoint}
  475. } else if t.shuffleReplicas {
  476. replicas = shuffleHosts(replicas)
  477. }
  478. var (
  479. fallbackIter NextHost
  480. i int
  481. )
  482. used := make(map[*HostInfo]bool, len(replicas))
  483. return func() SelectedHost {
  484. for i < len(replicas) {
  485. h := replicas[i]
  486. i++
  487. if h.IsUp() && t.fallback.IsLocal(h) {
  488. used[h] = true
  489. return (*selectedHost)(h)
  490. }
  491. }
  492. if fallbackIter == nil {
  493. // fallback
  494. fallbackIter = t.fallback.Pick(qry)
  495. }
  496. // filter the token aware selected hosts from the fallback hosts
  497. for fallbackHost := fallbackIter(); fallbackHost != nil; fallbackHost = fallbackIter() {
  498. if !used[fallbackHost.Info()] {
  499. return fallbackHost
  500. }
  501. }
  502. return nil
  503. }
  504. }
  505. // HostPoolHostPolicy is a host policy which uses the bitly/go-hostpool library
  506. // to distribute queries between hosts and prevent sending queries to
  507. // unresponsive hosts. When creating the host pool that is passed to the policy
  508. // use an empty slice of hosts as the hostpool will be populated later by gocql.
  509. // See below for examples of usage:
  510. //
  511. // // Create host selection policy using a simple host pool
  512. // cluster.PoolConfig.HostSelectionPolicy = HostPoolHostPolicy(hostpool.New(nil))
  513. //
  514. // // Create host selection policy using an epsilon greedy pool
  515. // cluster.PoolConfig.HostSelectionPolicy = HostPoolHostPolicy(
  516. // hostpool.NewEpsilonGreedy(nil, 0, &hostpool.LinearEpsilonValueCalculator{}),
  517. // )
  518. //
  519. func HostPoolHostPolicy(hp hostpool.HostPool) HostSelectionPolicy {
  520. return &hostPoolHostPolicy{hostMap: map[string]*HostInfo{}, hp: hp}
  521. }
  522. type hostPoolHostPolicy struct {
  523. hp hostpool.HostPool
  524. mu sync.RWMutex
  525. hostMap map[string]*HostInfo
  526. }
  527. func (r *hostPoolHostPolicy) Init(*Session) {}
  528. func (r *hostPoolHostPolicy) KeyspaceChanged(KeyspaceUpdateEvent) {}
  529. func (r *hostPoolHostPolicy) SetPartitioner(string) {}
  530. func (r *hostPoolHostPolicy) IsLocal(*HostInfo) bool { return true }
  531. func (r *hostPoolHostPolicy) SetHosts(hosts []*HostInfo) {
  532. peers := make([]string, len(hosts))
  533. hostMap := make(map[string]*HostInfo, len(hosts))
  534. for i, host := range hosts {
  535. ip := host.ConnectAddress().String()
  536. peers[i] = ip
  537. hostMap[ip] = host
  538. }
  539. r.mu.Lock()
  540. r.hp.SetHosts(peers)
  541. r.hostMap = hostMap
  542. r.mu.Unlock()
  543. }
  544. func (r *hostPoolHostPolicy) AddHost(host *HostInfo) {
  545. ip := host.ConnectAddress().String()
  546. r.mu.Lock()
  547. defer r.mu.Unlock()
  548. // If the host addr is present and isn't nil return
  549. if h, ok := r.hostMap[ip]; ok && h != nil {
  550. return
  551. }
  552. // otherwise, add the host to the map
  553. r.hostMap[ip] = host
  554. // and construct a new peer list to give to the HostPool
  555. hosts := make([]string, 0, len(r.hostMap))
  556. for addr := range r.hostMap {
  557. hosts = append(hosts, addr)
  558. }
  559. r.hp.SetHosts(hosts)
  560. }
  561. func (r *hostPoolHostPolicy) RemoveHost(host *HostInfo) {
  562. ip := host.ConnectAddress().String()
  563. r.mu.Lock()
  564. defer r.mu.Unlock()
  565. if _, ok := r.hostMap[ip]; !ok {
  566. return
  567. }
  568. delete(r.hostMap, ip)
  569. hosts := make([]string, 0, len(r.hostMap))
  570. for _, host := range r.hostMap {
  571. hosts = append(hosts, host.ConnectAddress().String())
  572. }
  573. r.hp.SetHosts(hosts)
  574. }
  575. func (r *hostPoolHostPolicy) HostUp(host *HostInfo) {
  576. r.AddHost(host)
  577. }
  578. func (r *hostPoolHostPolicy) HostDown(host *HostInfo) {
  579. r.RemoveHost(host)
  580. }
  581. func (r *hostPoolHostPolicy) Pick(qry ExecutableQuery) NextHost {
  582. return func() SelectedHost {
  583. r.mu.RLock()
  584. defer r.mu.RUnlock()
  585. if len(r.hostMap) == 0 {
  586. return nil
  587. }
  588. hostR := r.hp.Get()
  589. host, ok := r.hostMap[hostR.Host()]
  590. if !ok {
  591. return nil
  592. }
  593. return selectedHostPoolHost{
  594. policy: r,
  595. info: host,
  596. hostR: hostR,
  597. }
  598. }
  599. }
  600. // selectedHostPoolHost is a host returned by the hostPoolHostPolicy and
  601. // implements the SelectedHost interface
  602. type selectedHostPoolHost struct {
  603. policy *hostPoolHostPolicy
  604. info *HostInfo
  605. hostR hostpool.HostPoolResponse
  606. }
  607. func (host selectedHostPoolHost) Info() *HostInfo {
  608. return host.info
  609. }
  610. func (host selectedHostPoolHost) Mark(err error) {
  611. ip := host.info.ConnectAddress().String()
  612. host.policy.mu.RLock()
  613. defer host.policy.mu.RUnlock()
  614. if _, ok := host.policy.hostMap[ip]; !ok {
  615. // host was removed between pick and mark
  616. return
  617. }
  618. host.hostR.Mark(err)
  619. }
  620. type dcAwareRR struct {
  621. local string
  622. pos uint32
  623. mu sync.RWMutex
  624. localHosts cowHostList
  625. remoteHosts cowHostList
  626. }
  627. // DCAwareRoundRobinPolicy is a host selection policies which will prioritize and
  628. // return hosts which are in the local datacentre before returning hosts in all
  629. // other datercentres
  630. func DCAwareRoundRobinPolicy(localDC string) HostSelectionPolicy {
  631. return &dcAwareRR{local: localDC}
  632. }
  633. func (d *dcAwareRR) Init(*Session) {}
  634. func (d *dcAwareRR) KeyspaceChanged(KeyspaceUpdateEvent) {}
  635. func (d *dcAwareRR) SetPartitioner(p string) {}
  636. func (d *dcAwareRR) IsLocal(host *HostInfo) bool {
  637. return host.DataCenter() == d.local
  638. }
  639. func (d *dcAwareRR) AddHost(host *HostInfo) {
  640. if host.DataCenter() == d.local {
  641. d.localHosts.add(host)
  642. } else {
  643. d.remoteHosts.add(host)
  644. }
  645. }
  646. func (d *dcAwareRR) RemoveHost(host *HostInfo) {
  647. if host.DataCenter() == d.local {
  648. d.localHosts.remove(host.ConnectAddress())
  649. } else {
  650. d.remoteHosts.remove(host.ConnectAddress())
  651. }
  652. }
  653. func (d *dcAwareRR) HostUp(host *HostInfo) { d.AddHost(host) }
  654. func (d *dcAwareRR) HostDown(host *HostInfo) { d.RemoveHost(host) }
  655. func (d *dcAwareRR) Pick(q ExecutableQuery) NextHost {
  656. var i int
  657. return func() SelectedHost {
  658. var hosts []*HostInfo
  659. localHosts := d.localHosts.get()
  660. remoteHosts := d.remoteHosts.get()
  661. if len(localHosts) != 0 {
  662. hosts = localHosts
  663. } else {
  664. hosts = remoteHosts
  665. }
  666. if len(hosts) == 0 {
  667. return nil
  668. }
  669. // always increment pos to evenly distribute traffic in case of
  670. // failures
  671. pos := atomic.AddUint32(&d.pos, 1) - 1
  672. if i >= len(localHosts)+len(remoteHosts) {
  673. return nil
  674. }
  675. host := hosts[(pos)%uint32(len(hosts))]
  676. i++
  677. return (*selectedHost)(host)
  678. }
  679. }
  680. // ConvictionPolicy interface is used by gocql to determine if a host should be
  681. // marked as DOWN based on the error and host info
  682. type ConvictionPolicy interface {
  683. // Implementations should return `true` if the host should be convicted, `false` otherwise.
  684. AddFailure(error error, host *HostInfo) bool
  685. //Implementations should clear out any convictions or state regarding the host.
  686. Reset(host *HostInfo)
  687. }
  688. // SimpleConvictionPolicy implements a ConvictionPolicy which convicts all hosts
  689. // regardless of error
  690. type SimpleConvictionPolicy struct {
  691. }
  692. func (e *SimpleConvictionPolicy) AddFailure(error error, host *HostInfo) bool {
  693. return true
  694. }
  695. func (e *SimpleConvictionPolicy) Reset(host *HostInfo) {}
  696. // ReconnectionPolicy interface is used by gocql to determine if reconnection
  697. // can be attempted after connection error. The interface allows gocql users
  698. // to implement their own logic to determine how to attempt reconnection.
  699. //
  700. type ReconnectionPolicy interface {
  701. GetInterval(currentRetry int) time.Duration
  702. GetMaxRetries() int
  703. }
  704. // ConstantReconnectionPolicy has simple logic for returning a fixed reconnection interval.
  705. //
  706. // Examples of usage:
  707. //
  708. // cluster.ReconnectionPolicy = &gocql.ConstantReconnectionPolicy{MaxRetries: 10, Interval: 8 * time.Second}
  709. //
  710. type ConstantReconnectionPolicy struct {
  711. MaxRetries int
  712. Interval time.Duration
  713. }
  714. func (c *ConstantReconnectionPolicy) GetInterval(currentRetry int) time.Duration {
  715. return c.Interval
  716. }
  717. func (c *ConstantReconnectionPolicy) GetMaxRetries() int {
  718. return c.MaxRetries
  719. }
  720. // ExponentialReconnectionPolicy returns a growing reconnection interval.
  721. type ExponentialReconnectionPolicy struct {
  722. MaxRetries int
  723. InitialInterval time.Duration
  724. }
  725. func (e *ExponentialReconnectionPolicy) GetInterval(currentRetry int) time.Duration {
  726. return getExponentialTime(e.InitialInterval, math.MaxInt16*time.Second, e.GetMaxRetries())
  727. }
  728. func (e *ExponentialReconnectionPolicy) GetMaxRetries() int {
  729. return e.MaxRetries
  730. }
  731. type SpeculativeExecutionPolicy interface {
  732. Attempts() int
  733. Delay() time.Duration
  734. }
  735. type NonSpeculativeExecution struct{}
  736. func (sp NonSpeculativeExecution) Attempts() int { return 0 } // No additional attempts
  737. func (sp NonSpeculativeExecution) Delay() time.Duration { return 1 } // The delay. Must be positive to be used in a ticker.
  738. type SimpleSpeculativeExecution struct {
  739. NumAttempts int
  740. TimeoutDelay time.Duration
  741. }
  742. func (sp *SimpleSpeculativeExecution) Attempts() int { return sp.NumAttempts }
  743. func (sp *SimpleSpeculativeExecution) Delay() time.Duration { return sp.TimeoutDelay }