conn.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014
  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. package gocql
  5. import (
  6. "bufio"
  7. "crypto/tls"
  8. "errors"
  9. "fmt"
  10. "io"
  11. "io/ioutil"
  12. "log"
  13. "net"
  14. "strconv"
  15. "strings"
  16. "sync"
  17. "sync/atomic"
  18. "time"
  19. )
  20. var (
  21. approvedAuthenticators = [...]string{
  22. "org.apache.cassandra.auth.PasswordAuthenticator",
  23. "com.instaclustr.cassandra.auth.SharedSecretAuthenticator",
  24. }
  25. )
  26. func approve(authenticator string) bool {
  27. for _, s := range approvedAuthenticators {
  28. if authenticator == s {
  29. return true
  30. }
  31. }
  32. return false
  33. }
  34. //JoinHostPort is a utility to return a address string that can be used
  35. //gocql.Conn to form a connection with a host.
  36. func JoinHostPort(addr string, port int) string {
  37. addr = strings.TrimSpace(addr)
  38. if _, _, err := net.SplitHostPort(addr); err != nil {
  39. addr = net.JoinHostPort(addr, strconv.Itoa(port))
  40. }
  41. return addr
  42. }
  43. type Authenticator interface {
  44. Challenge(req []byte) (resp []byte, auth Authenticator, err error)
  45. Success(data []byte) error
  46. }
  47. type PasswordAuthenticator struct {
  48. Username string
  49. Password string
  50. }
  51. func (p PasswordAuthenticator) Challenge(req []byte) ([]byte, Authenticator, error) {
  52. if !approve(string(req)) {
  53. return nil, nil, fmt.Errorf("unexpected authenticator %q", req)
  54. }
  55. resp := make([]byte, 2+len(p.Username)+len(p.Password))
  56. resp[0] = 0
  57. copy(resp[1:], p.Username)
  58. resp[len(p.Username)+1] = 0
  59. copy(resp[2+len(p.Username):], p.Password)
  60. return resp, nil, nil
  61. }
  62. func (p PasswordAuthenticator) Success(data []byte) error {
  63. return nil
  64. }
  65. type SslOptions struct {
  66. tls.Config
  67. // CertPath and KeyPath are optional depending on server
  68. // config, but both fields must be omitted to avoid using a
  69. // client certificate
  70. CertPath string
  71. KeyPath string
  72. CaPath string //optional depending on server config
  73. // If you want to verify the hostname and server cert (like a wildcard for cass cluster) then you should turn this on
  74. // This option is basically the inverse of InSecureSkipVerify
  75. // See InSecureSkipVerify in http://golang.org/pkg/crypto/tls/ for more info
  76. EnableHostVerification bool
  77. }
  78. type ConnConfig struct {
  79. ProtoVersion int
  80. CQLVersion string
  81. Timeout time.Duration
  82. NumStreams int
  83. Compressor Compressor
  84. Authenticator Authenticator
  85. Keepalive time.Duration
  86. tlsConfig *tls.Config
  87. }
  88. type ConnErrorHandler interface {
  89. HandleError(conn *Conn, err error, closed bool)
  90. }
  91. // How many timeouts we will allow to occur before the connection is closed
  92. // and restarted. This is to prevent a single query timeout from killing a connection
  93. // which may be serving more queries just fine.
  94. // Default is 10, should not be changed concurrently with queries.
  95. var TimeoutLimit int64 = 10
  96. // Conn is a single connection to a Cassandra node. It can be used to execute
  97. // queries, but users are usually advised to use a more reliable, higher
  98. // level API.
  99. type Conn struct {
  100. conn net.Conn
  101. r *bufio.Reader
  102. timeout time.Duration
  103. cfg *ConnConfig
  104. numStreams int
  105. headerBuf []byte
  106. uniq chan int
  107. mu sync.RWMutex
  108. calls map[int]*callReq
  109. errorHandler ConnErrorHandler
  110. compressor Compressor
  111. auth Authenticator
  112. addr string
  113. version uint8
  114. currentKeyspace string
  115. started bool
  116. session *Session
  117. closed int32
  118. quit chan struct{}
  119. timeouts int64
  120. }
  121. // Connect establishes a connection to a Cassandra node.
  122. // You must also call the Serve method before you can execute any queries.
  123. func Connect(addr string, cfg *ConnConfig, errorHandler ConnErrorHandler, session *Session) (*Conn, error) {
  124. var (
  125. err error
  126. conn net.Conn
  127. )
  128. dialer := &net.Dialer{
  129. Timeout: cfg.Timeout,
  130. }
  131. if cfg.tlsConfig != nil {
  132. // the TLS config is safe to be reused by connections but it must not
  133. // be modified after being used.
  134. conn, err = tls.DialWithDialer(dialer, "tcp", addr, cfg.tlsConfig)
  135. } else {
  136. conn, err = dialer.Dial("tcp", addr)
  137. }
  138. if err != nil {
  139. return nil, err
  140. }
  141. // going to default to proto 2
  142. if cfg.ProtoVersion < protoVersion1 || cfg.ProtoVersion > protoVersion4 {
  143. log.Printf("unsupported protocol version: %d using 2\n", cfg.ProtoVersion)
  144. cfg.ProtoVersion = 2
  145. }
  146. headerSize := 8
  147. maxStreams := 128
  148. if cfg.ProtoVersion > protoVersion2 {
  149. maxStreams = 32768
  150. headerSize = 9
  151. }
  152. streams := cfg.NumStreams
  153. if streams <= 0 || streams >= maxStreams {
  154. streams = maxStreams
  155. } else {
  156. streams++
  157. }
  158. c := &Conn{
  159. conn: conn,
  160. r: bufio.NewReader(conn),
  161. cfg: cfg,
  162. uniq: make(chan int, streams),
  163. calls: make(map[int]*callReq),
  164. timeout: cfg.Timeout,
  165. version: uint8(cfg.ProtoVersion),
  166. addr: conn.RemoteAddr().String(),
  167. errorHandler: errorHandler,
  168. compressor: cfg.Compressor,
  169. auth: cfg.Authenticator,
  170. headerBuf: make([]byte, headerSize),
  171. quit: make(chan struct{}),
  172. session: session,
  173. numStreams: streams,
  174. }
  175. if cfg.Keepalive > 0 {
  176. c.setKeepalive(cfg.Keepalive)
  177. }
  178. // reserve stream 0 incase cassandra returns an error on it without us sending
  179. // a request.
  180. for i := 1; i < streams; i++ {
  181. c.uniq <- i
  182. }
  183. go c.serve()
  184. if err := c.startup(); err != nil {
  185. conn.Close()
  186. return nil, err
  187. }
  188. c.started = true
  189. return c, nil
  190. }
  191. func (c *Conn) Write(p []byte) (int, error) {
  192. if c.timeout > 0 {
  193. c.conn.SetWriteDeadline(time.Now().Add(c.timeout))
  194. }
  195. return c.conn.Write(p)
  196. }
  197. func (c *Conn) Read(p []byte) (n int, err error) {
  198. const maxAttempts = 5
  199. for i := 0; i < maxAttempts; i++ {
  200. var nn int
  201. if c.timeout > 0 {
  202. c.conn.SetReadDeadline(time.Now().Add(c.timeout))
  203. }
  204. nn, err = io.ReadFull(c.r, p[n:])
  205. n += nn
  206. if err == nil {
  207. break
  208. }
  209. if verr, ok := err.(net.Error); !ok || !verr.Temporary() {
  210. break
  211. }
  212. }
  213. return
  214. }
  215. func (c *Conn) startup() error {
  216. m := map[string]string{
  217. "CQL_VERSION": c.cfg.CQLVersion,
  218. }
  219. if c.compressor != nil {
  220. m["COMPRESSION"] = c.compressor.Name()
  221. }
  222. framer, err := c.exec(&writeStartupFrame{opts: m}, nil)
  223. if err != nil {
  224. return err
  225. }
  226. frame, err := framer.parseFrame()
  227. if err != nil {
  228. return err
  229. }
  230. switch v := frame.(type) {
  231. case error:
  232. return v
  233. case *readyFrame:
  234. return nil
  235. case *authenticateFrame:
  236. return c.authenticateHandshake(v)
  237. default:
  238. return NewErrProtocol("Unknown type of response to startup frame: %s", v)
  239. }
  240. }
  241. func (c *Conn) authenticateHandshake(authFrame *authenticateFrame) error {
  242. if c.auth == nil {
  243. return fmt.Errorf("authentication required (using %q)", authFrame.class)
  244. }
  245. resp, challenger, err := c.auth.Challenge([]byte(authFrame.class))
  246. if err != nil {
  247. return err
  248. }
  249. req := &writeAuthResponseFrame{data: resp}
  250. for {
  251. framer, err := c.exec(req, nil)
  252. if err != nil {
  253. return err
  254. }
  255. frame, err := framer.parseFrame()
  256. if err != nil {
  257. return err
  258. }
  259. switch v := frame.(type) {
  260. case error:
  261. return v
  262. case *authSuccessFrame:
  263. if challenger != nil {
  264. return challenger.Success(v.data)
  265. }
  266. return nil
  267. case *authChallengeFrame:
  268. resp, challenger, err = challenger.Challenge(v.data)
  269. if err != nil {
  270. return err
  271. }
  272. req = &writeAuthResponseFrame{
  273. data: resp,
  274. }
  275. default:
  276. return fmt.Errorf("unknown frame response during authentication: %v", v)
  277. }
  278. framerPool.Put(framer)
  279. }
  280. }
  281. func (c *Conn) closeWithError(err error) {
  282. if !atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
  283. return
  284. }
  285. if err != nil {
  286. // we should attempt to deliver the error back to the caller if it
  287. // exists
  288. c.mu.RLock()
  289. for _, req := range c.calls {
  290. // we need to send the error to all waiting queries, put the state
  291. // of this conn into not active so that it can not execute any queries.
  292. if err != nil {
  293. select {
  294. case req.resp <- err:
  295. default:
  296. }
  297. }
  298. }
  299. c.mu.RUnlock()
  300. }
  301. // if error was nil then unblock the quit channel
  302. close(c.quit)
  303. c.conn.Close()
  304. if c.started && err != nil {
  305. c.errorHandler.HandleError(c, err, true)
  306. }
  307. }
  308. func (c *Conn) Close() {
  309. c.closeWithError(nil)
  310. }
  311. // Serve starts the stream multiplexer for this connection, which is required
  312. // to execute any queries. This method runs as long as the connection is
  313. // open and is therefore usually called in a separate goroutine.
  314. func (c *Conn) serve() {
  315. var (
  316. err error
  317. )
  318. for {
  319. err = c.recv()
  320. if err != nil {
  321. break
  322. }
  323. }
  324. c.closeWithError(err)
  325. }
  326. func (c *Conn) discardFrame(head frameHeader) error {
  327. _, err := io.CopyN(ioutil.Discard, c, int64(head.length))
  328. if err != nil {
  329. return err
  330. }
  331. return nil
  332. }
  333. func (c *Conn) recv() error {
  334. // not safe for concurrent reads
  335. // read a full header, ignore timeouts, as this is being ran in a loop
  336. // TODO: TCP level deadlines? or just query level deadlines?
  337. if c.timeout > 0 {
  338. c.conn.SetReadDeadline(time.Time{})
  339. }
  340. // were just reading headers over and over and copy bodies
  341. head, err := readHeader(c.r, c.headerBuf)
  342. if err != nil {
  343. return err
  344. }
  345. if head.stream > c.numStreams {
  346. return fmt.Errorf("gocql: frame header stream is beyond call exepected bounds: %d", head.stream)
  347. } else if head.stream == -1 {
  348. // TODO: handle cassandra event frames, we shouldnt get any currently
  349. return c.discardFrame(head)
  350. } else if head.stream <= 0 {
  351. // reserved stream that we dont use, probably due to a protocol error
  352. // or a bug in Cassandra, this should be an error, parse it and return.
  353. framer := newFramer(c, c, c.compressor, c.version)
  354. if err := framer.readFrame(&head); err != nil {
  355. return err
  356. }
  357. defer framerPool.Put(framer)
  358. frame, err := framer.parseFrame()
  359. if err != nil {
  360. return err
  361. }
  362. switch v := frame.(type) {
  363. case error:
  364. return fmt.Errorf("gocql: error on stream %d: %v", head.stream, v)
  365. default:
  366. return fmt.Errorf("gocql: received frame on stream %d: %v", head.stream, frame)
  367. }
  368. }
  369. c.mu.RLock()
  370. call, ok := c.calls[head.stream]
  371. c.mu.RUnlock()
  372. if call == nil || call.framer == nil || !ok {
  373. log.Printf("gocql: received response for stream which has no handler: header=%v\n", head)
  374. return c.discardFrame(head)
  375. }
  376. err = call.framer.readFrame(&head)
  377. if err != nil {
  378. // only net errors should cause the connection to be closed. Though
  379. // cassandra returning corrupt frames will be returned here as well.
  380. if _, ok := err.(net.Error); ok {
  381. return err
  382. }
  383. }
  384. // we either, return a response to the caller, the caller timedout, or the
  385. // connection has closed. Either way we should never block indefinatly here
  386. select {
  387. case call.resp <- err:
  388. case <-call.timeout:
  389. c.releaseStream(head.stream)
  390. case <-c.quit:
  391. }
  392. return nil
  393. }
  394. type callReq struct {
  395. // could use a waitgroup but this allows us to do timeouts on the read/send
  396. resp chan error
  397. framer *framer
  398. timeout chan struct{} // indicates to recv() that a call has timedout
  399. streamID int // current stream in use
  400. }
  401. func (c *Conn) releaseStream(stream int) {
  402. c.mu.Lock()
  403. call := c.calls[stream]
  404. if call != nil && stream != call.streamID {
  405. panic(fmt.Sprintf("attempt to release streamID with ivalid stream: %d -> %+v\n", stream, call))
  406. }
  407. delete(c.calls, stream)
  408. c.mu.Unlock()
  409. streamPool.Put(call)
  410. select {
  411. case c.uniq <- stream:
  412. case <-c.quit:
  413. }
  414. }
  415. func (c *Conn) handleTimeout() {
  416. if atomic.AddInt64(&c.timeouts, 1) > TimeoutLimit {
  417. c.closeWithError(ErrTooManyTimeouts)
  418. }
  419. }
  420. var (
  421. streamPool = sync.Pool{
  422. New: func() interface{} {
  423. return &callReq{
  424. resp: make(chan error),
  425. }
  426. },
  427. }
  428. )
  429. func (c *Conn) exec(req frameWriter, tracer Tracer) (*framer, error) {
  430. // TODO: move tracer onto conn
  431. var stream int
  432. select {
  433. case stream = <-c.uniq:
  434. case <-c.quit:
  435. return nil, ErrConnectionClosed
  436. }
  437. // resp is basically a waiting semaphore protecting the framer
  438. framer := newFramer(c, c, c.compressor, c.version)
  439. c.mu.Lock()
  440. call := c.calls[stream]
  441. if call != nil {
  442. panic(fmt.Sprintf("attempting to use stream already in use: %d -> %+v\n", stream, call))
  443. } else {
  444. call = streamPool.Get().(*callReq)
  445. }
  446. c.calls[stream] = call
  447. c.mu.Unlock()
  448. call.framer = framer
  449. call.timeout = make(chan struct{})
  450. call.streamID = stream
  451. if tracer != nil {
  452. framer.trace()
  453. }
  454. err := req.writeFrame(framer, stream)
  455. if err != nil {
  456. // I think this is the correct thing to do, im not entirely sure. It is not
  457. // ideal as readers might still get some data, but they probably wont.
  458. // Here we need to be careful as the stream is not available and if all
  459. // writes just timeout or fail then the pool might use this connection to
  460. // send a frame on, with all the streams used up and not returned.
  461. c.closeWithError(err)
  462. return nil, err
  463. }
  464. select {
  465. case err := <-call.resp:
  466. if err != nil {
  467. if !c.Closed() {
  468. // if the connection is closed then we cant release the stream,
  469. // this is because the request is still outstanding and we have
  470. // been handed another error from another stream which caused the
  471. // connection to close.
  472. c.releaseStream(stream)
  473. }
  474. return nil, err
  475. }
  476. case <-time.After(c.timeout):
  477. close(call.timeout)
  478. c.handleTimeout()
  479. return nil, ErrTimeoutNoResponse
  480. case <-c.quit:
  481. return nil, ErrConnectionClosed
  482. }
  483. // dont release the stream if detect a timeout as another request can reuse
  484. // that stream and get a response for the old request, which we have no
  485. // easy way of detecting.
  486. //
  487. // Ensure that the stream is not released if there are potentially outstanding
  488. // requests on the stream to prevent nil pointer dereferences in recv().
  489. defer c.releaseStream(stream)
  490. if v := framer.header.version.version(); v != c.version {
  491. return nil, NewErrProtocol("unexpected protocol version in response: got %d expected %d", v, c.version)
  492. }
  493. return framer, nil
  494. }
  495. func (c *Conn) prepareStatement(stmt string, tracer Tracer) (*QueryInfo, error) {
  496. stmtsLRU.Lock()
  497. if stmtsLRU.lru == nil {
  498. initStmtsLRU(defaultMaxPreparedStmts)
  499. }
  500. stmtCacheKey := c.addr + c.currentKeyspace + stmt
  501. if val, ok := stmtsLRU.lru.Get(stmtCacheKey); ok {
  502. stmtsLRU.Unlock()
  503. flight := val.(*inflightPrepare)
  504. flight.wg.Wait()
  505. return &flight.info, flight.err
  506. }
  507. flight := new(inflightPrepare)
  508. flight.wg.Add(1)
  509. stmtsLRU.lru.Add(stmtCacheKey, flight)
  510. stmtsLRU.Unlock()
  511. prep := &writePrepareFrame{
  512. statement: stmt,
  513. }
  514. framer, err := c.exec(prep, tracer)
  515. if err != nil {
  516. flight.err = err
  517. flight.wg.Done()
  518. return nil, err
  519. }
  520. frame, err := framer.parseFrame()
  521. if err != nil {
  522. flight.err = err
  523. flight.wg.Done()
  524. return nil, err
  525. }
  526. // TODO(zariel): tidy this up, simplify handling of frame parsing so its not duplicated
  527. // everytime we need to parse a frame.
  528. if len(framer.traceID) > 0 {
  529. tracer.Trace(framer.traceID)
  530. }
  531. switch x := frame.(type) {
  532. case *resultPreparedFrame:
  533. // defensivly copy as we will recycle the underlying buffer after we
  534. // return.
  535. flight.info.Id = copyBytes(x.preparedID)
  536. // the type info's should _not_ have a reference to the framers read buffer,
  537. // therefore we can just copy them directly.
  538. flight.info.Args = x.reqMeta.columns
  539. flight.info.PKeyColumns = x.reqMeta.pkeyColumns
  540. flight.info.Rval = x.respMeta.columns
  541. case error:
  542. flight.err = x
  543. default:
  544. flight.err = NewErrProtocol("Unknown type in response to prepare frame: %s", x)
  545. }
  546. flight.wg.Done()
  547. if flight.err != nil {
  548. stmtsLRU.Lock()
  549. stmtsLRU.lru.Remove(stmtCacheKey)
  550. stmtsLRU.Unlock()
  551. }
  552. framerPool.Put(framer)
  553. return &flight.info, flight.err
  554. }
  555. func (c *Conn) executeQuery(qry *Query) *Iter {
  556. params := queryParams{
  557. consistency: qry.cons,
  558. }
  559. // frame checks that it is not 0
  560. params.serialConsistency = qry.serialCons
  561. params.defaultTimestamp = qry.defaultTimestamp
  562. if len(qry.pageState) > 0 {
  563. params.pagingState = qry.pageState
  564. }
  565. if qry.pageSize > 0 {
  566. params.pageSize = qry.pageSize
  567. }
  568. var frame frameWriter
  569. if qry.shouldPrepare() {
  570. // Prepare all DML queries. Other queries can not be prepared.
  571. info, err := c.prepareStatement(qry.stmt, qry.trace)
  572. if err != nil {
  573. return &Iter{err: err}
  574. }
  575. var values []interface{}
  576. if qry.binding == nil {
  577. values = qry.values
  578. } else {
  579. values, err = qry.binding(info)
  580. if err != nil {
  581. return &Iter{err: err}
  582. }
  583. }
  584. if len(values) != len(info.Args) {
  585. return &Iter{err: ErrQueryArgLength}
  586. }
  587. params.values = make([]queryValues, len(values))
  588. for i := 0; i < len(values); i++ {
  589. val, err := Marshal(info.Args[i].TypeInfo, values[i])
  590. if err != nil {
  591. return &Iter{err: err}
  592. }
  593. v := &params.values[i]
  594. v.value = val
  595. // TODO: handle query binding names
  596. }
  597. frame = &writeExecuteFrame{
  598. preparedID: info.Id,
  599. params: params,
  600. }
  601. } else {
  602. frame = &writeQueryFrame{
  603. statement: qry.stmt,
  604. params: params,
  605. }
  606. }
  607. framer, err := c.exec(frame, qry.trace)
  608. if err != nil {
  609. return &Iter{err: err}
  610. }
  611. resp, err := framer.parseFrame()
  612. if err != nil {
  613. return &Iter{err: err}
  614. }
  615. if len(framer.traceID) > 0 {
  616. qry.trace.Trace(framer.traceID)
  617. }
  618. switch x := resp.(type) {
  619. case *resultVoidFrame:
  620. return &Iter{framer: framer}
  621. case *resultRowsFrame:
  622. iter := &Iter{
  623. meta: x.meta,
  624. rows: x.rows,
  625. framer: framer,
  626. }
  627. if len(x.meta.pagingState) > 0 && !qry.disableAutoPage {
  628. iter.next = &nextIter{
  629. qry: *qry,
  630. pos: int((1 - qry.prefetch) * float64(len(iter.rows))),
  631. }
  632. iter.next.qry.pageState = x.meta.pagingState
  633. if iter.next.pos < 1 {
  634. iter.next.pos = 1
  635. }
  636. }
  637. return iter
  638. case *resultKeyspaceFrame:
  639. return &Iter{framer: framer}
  640. case *resultSchemaChangeFrame, *schemaChangeKeyspace, *schemaChangeTable, *schemaChangeFunction:
  641. iter := &Iter{framer: framer}
  642. c.awaitSchemaAgreement()
  643. // dont return an error from this, might be a good idea to give a warning
  644. // though. The impact of this returning an error would be that the cluster
  645. // is not consistent with regards to its schema.
  646. return iter
  647. case *RequestErrUnprepared:
  648. stmtsLRU.Lock()
  649. stmtCacheKey := c.addr + c.currentKeyspace + qry.stmt
  650. if _, ok := stmtsLRU.lru.Get(stmtCacheKey); ok {
  651. stmtsLRU.lru.Remove(stmtCacheKey)
  652. stmtsLRU.Unlock()
  653. return c.executeQuery(qry)
  654. }
  655. stmtsLRU.Unlock()
  656. return &Iter{err: x, framer: framer}
  657. case error:
  658. return &Iter{err: x, framer: framer}
  659. default:
  660. return &Iter{
  661. err: NewErrProtocol("Unknown type in response to execute query (%T): %s", x, x),
  662. framer: framer,
  663. }
  664. }
  665. }
  666. func (c *Conn) Pick(qry *Query) *Conn {
  667. if c.Closed() {
  668. return nil
  669. }
  670. return c
  671. }
  672. func (c *Conn) Closed() bool {
  673. return atomic.LoadInt32(&c.closed) == 1
  674. }
  675. func (c *Conn) Address() string {
  676. return c.addr
  677. }
  678. func (c *Conn) AvailableStreams() int {
  679. return len(c.uniq)
  680. }
  681. func (c *Conn) UseKeyspace(keyspace string) error {
  682. q := &writeQueryFrame{statement: `USE "` + keyspace + `"`}
  683. q.params.consistency = Any
  684. framer, err := c.exec(q, nil)
  685. if err != nil {
  686. return err
  687. }
  688. resp, err := framer.parseFrame()
  689. if err != nil {
  690. return err
  691. }
  692. switch x := resp.(type) {
  693. case *resultKeyspaceFrame:
  694. case error:
  695. return x
  696. default:
  697. return NewErrProtocol("unknown frame in response to USE: %v", x)
  698. }
  699. c.currentKeyspace = keyspace
  700. return nil
  701. }
  702. func (c *Conn) executeBatch(batch *Batch) (*Iter, error) {
  703. if c.version == protoVersion1 {
  704. return nil, ErrUnsupported
  705. }
  706. n := len(batch.Entries)
  707. req := &writeBatchFrame{
  708. typ: batch.Type,
  709. statements: make([]batchStatment, n),
  710. consistency: batch.Cons,
  711. serialConsistency: batch.serialCons,
  712. defaultTimestamp: batch.defaultTimestamp,
  713. }
  714. stmts := make(map[string]string)
  715. for i := 0; i < n; i++ {
  716. entry := &batch.Entries[i]
  717. b := &req.statements[i]
  718. if len(entry.Args) > 0 || entry.binding != nil {
  719. info, err := c.prepareStatement(entry.Stmt, nil)
  720. if err != nil {
  721. return nil, err
  722. }
  723. var args []interface{}
  724. if entry.binding == nil {
  725. args = entry.Args
  726. } else {
  727. args, err = entry.binding(info)
  728. if err != nil {
  729. return nil, err
  730. }
  731. }
  732. if len(args) != len(info.Args) {
  733. return nil, ErrQueryArgLength
  734. }
  735. b.preparedID = info.Id
  736. stmts[string(info.Id)] = entry.Stmt
  737. b.values = make([]queryValues, len(info.Args))
  738. for j := 0; j < len(info.Args); j++ {
  739. val, err := Marshal(info.Args[j].TypeInfo, args[j])
  740. if err != nil {
  741. return nil, err
  742. }
  743. b.values[j].value = val
  744. // TODO: add names
  745. }
  746. } else {
  747. b.statement = entry.Stmt
  748. }
  749. }
  750. // TODO: should batch support tracing?
  751. framer, err := c.exec(req, nil)
  752. if err != nil {
  753. return nil, err
  754. }
  755. resp, err := framer.parseFrame()
  756. if err != nil {
  757. return nil, err
  758. }
  759. switch x := resp.(type) {
  760. case *resultVoidFrame:
  761. framerPool.Put(framer)
  762. return nil, nil
  763. case *RequestErrUnprepared:
  764. stmt, found := stmts[string(x.StatementId)]
  765. if found {
  766. stmtsLRU.Lock()
  767. stmtsLRU.lru.Remove(c.addr + c.currentKeyspace + stmt)
  768. stmtsLRU.Unlock()
  769. }
  770. framerPool.Put(framer)
  771. if found {
  772. return c.executeBatch(batch)
  773. } else {
  774. return nil, x
  775. }
  776. case *resultRowsFrame:
  777. iter := &Iter{
  778. meta: x.meta,
  779. rows: x.rows,
  780. framer: framer,
  781. }
  782. return iter, nil
  783. case error:
  784. framerPool.Put(framer)
  785. return nil, x
  786. default:
  787. framerPool.Put(framer)
  788. return nil, NewErrProtocol("Unknown type in response to batch statement: %s", x)
  789. }
  790. }
  791. func (c *Conn) setKeepalive(d time.Duration) error {
  792. if tc, ok := c.conn.(*net.TCPConn); ok {
  793. err := tc.SetKeepAlivePeriod(d)
  794. if err != nil {
  795. return err
  796. }
  797. return tc.SetKeepAlive(true)
  798. }
  799. return nil
  800. }
  801. func (c *Conn) query(statement string, values ...interface{}) (iter *Iter) {
  802. q := c.session.Query(statement, values...).Consistency(One)
  803. return c.executeQuery(q)
  804. }
  805. func (c *Conn) awaitSchemaAgreement() (err error) {
  806. const (
  807. peerSchemas = "SELECT schema_version FROM system.peers"
  808. localSchemas = "SELECT schema_version FROM system.local WHERE key='local'"
  809. )
  810. endDeadline := time.Now().Add(c.session.cfg.MaxWaitSchemaAgreement)
  811. for time.Now().Before(endDeadline) {
  812. iter := c.query(peerSchemas)
  813. versions := make(map[string]struct{})
  814. var schemaVersion string
  815. for iter.Scan(&schemaVersion) {
  816. versions[schemaVersion] = struct{}{}
  817. schemaVersion = ""
  818. }
  819. if err = iter.Close(); err != nil {
  820. goto cont
  821. }
  822. iter = c.query(localSchemas)
  823. for iter.Scan(&schemaVersion) {
  824. versions[schemaVersion] = struct{}{}
  825. schemaVersion = ""
  826. }
  827. if err = iter.Close(); err != nil {
  828. goto cont
  829. }
  830. if len(versions) <= 1 {
  831. return nil
  832. }
  833. cont:
  834. time.Sleep(200 * time.Millisecond)
  835. }
  836. if err != nil {
  837. return
  838. }
  839. // not exported
  840. return errors.New("gocql: cluster schema versions not consistent")
  841. }
  842. type inflightPrepare struct {
  843. info QueryInfo
  844. err error
  845. wg sync.WaitGroup
  846. }
  847. var (
  848. ErrQueryArgLength = errors.New("gocql: query argument length mismatch")
  849. ErrTimeoutNoResponse = errors.New("gocql: no response received from cassandra within timeout period")
  850. ErrTooManyTimeouts = errors.New("gocql: too many query timeouts on the connection")
  851. ErrConnectionClosed = errors.New("gocql: connection closed waiting for response")
  852. )