conn.go 29 KB

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