conn.go 28 KB

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