policies.go 24 KB

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