conn.go 31 KB

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