conn.go 27 KB

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