conn.go 26 KB

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