conn.go 32 KB

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