conn.go 26 KB

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