conn.go 23 KB

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