conn.go 27 KB

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