conn.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798
  1. // Copyright 2012 Gary Burd
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License"): you may
  4. // not use this file except in compliance with the License. You may obtain
  5. // a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  11. // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  12. // License for the specific language governing permissions and limitations
  13. // under the License.
  14. package redis
  15. import (
  16. "bufio"
  17. "bytes"
  18. "context"
  19. "crypto/tls"
  20. "errors"
  21. "fmt"
  22. "io"
  23. "net"
  24. "net/url"
  25. "regexp"
  26. "strconv"
  27. "sync"
  28. "time"
  29. )
  30. var (
  31. _ ConnWithTimeout = (*conn)(nil)
  32. )
  33. // conn is the low-level implementation of Conn
  34. type conn struct {
  35. // Shared
  36. mu sync.Mutex
  37. pending int
  38. err error
  39. conn net.Conn
  40. // Read
  41. readTimeout time.Duration
  42. br *bufio.Reader
  43. // Write
  44. writeTimeout time.Duration
  45. bw *bufio.Writer
  46. // Scratch space for formatting argument length.
  47. // '*' or '$', length, "\r\n"
  48. lenScratch [32]byte
  49. // Scratch space for formatting integers and floats.
  50. numScratch [40]byte
  51. }
  52. // DialTimeout acts like Dial but takes timeouts for establishing the
  53. // connection to the server, writing a command and reading a reply.
  54. //
  55. // Deprecated: Use Dial with options instead.
  56. func DialTimeout(network, address string, connectTimeout, readTimeout, writeTimeout time.Duration) (Conn, error) {
  57. return Dial(network, address,
  58. DialConnectTimeout(connectTimeout),
  59. DialReadTimeout(readTimeout),
  60. DialWriteTimeout(writeTimeout))
  61. }
  62. // DialOption specifies an option for dialing a Redis server.
  63. type DialOption struct {
  64. f func(*dialOptions)
  65. }
  66. type dialOptions struct {
  67. readTimeout time.Duration
  68. writeTimeout time.Duration
  69. tlsHandshakeTimeout time.Duration
  70. dialer *net.Dialer
  71. dialContext func(ctx context.Context, network, addr string) (net.Conn, error)
  72. db int
  73. username string
  74. password string
  75. clientName string
  76. useTLS bool
  77. skipVerify bool
  78. tlsConfig *tls.Config
  79. }
  80. // DialTLSHandshakeTimeout specifies the maximum amount of time waiting to
  81. // wait for a TLS handshake. Zero means no timeout.
  82. // If no DialTLSHandshakeTimeout option is specified then the default is 30 seconds.
  83. func DialTLSHandshakeTimeout(d time.Duration) DialOption {
  84. return DialOption{func(do *dialOptions) {
  85. do.tlsHandshakeTimeout = d
  86. }}
  87. }
  88. // DialReadTimeout specifies the timeout for reading a single command reply.
  89. func DialReadTimeout(d time.Duration) DialOption {
  90. return DialOption{func(do *dialOptions) {
  91. do.readTimeout = d
  92. }}
  93. }
  94. // DialWriteTimeout specifies the timeout for writing a single command.
  95. func DialWriteTimeout(d time.Duration) DialOption {
  96. return DialOption{func(do *dialOptions) {
  97. do.writeTimeout = d
  98. }}
  99. }
  100. // DialConnectTimeout specifies the timeout for connecting to the Redis server when
  101. // no DialNetDial option is specified.
  102. // If no DialConnectTimeout option is specified then the default is 30 seconds.
  103. func DialConnectTimeout(d time.Duration) DialOption {
  104. return DialOption{func(do *dialOptions) {
  105. do.dialer.Timeout = d
  106. }}
  107. }
  108. // DialKeepAlive specifies the keep-alive period for TCP connections to the Redis server
  109. // when no DialNetDial option is specified.
  110. // If zero, keep-alives are not enabled. If no DialKeepAlive option is specified then
  111. // the default of 5 minutes is used to ensure that half-closed TCP sessions are detected.
  112. func DialKeepAlive(d time.Duration) DialOption {
  113. return DialOption{func(do *dialOptions) {
  114. do.dialer.KeepAlive = d
  115. }}
  116. }
  117. // DialNetDial specifies a custom dial function for creating TCP
  118. // connections, otherwise a net.Dialer customized via the other options is used.
  119. // DialNetDial overrides DialConnectTimeout and DialKeepAlive.
  120. func DialNetDial(dial func(network, addr string) (net.Conn, error)) DialOption {
  121. return DialOption{func(do *dialOptions) {
  122. do.dialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
  123. return dial(network, addr)
  124. }
  125. }}
  126. }
  127. // DialContextFunc specifies a custom dial function with context for creating TCP
  128. // connections, otherwise a net.Dialer customized via the other options is used.
  129. // DialContextFunc overrides DialConnectTimeout and DialKeepAlive.
  130. func DialContextFunc(f func(ctx context.Context, network, addr string) (net.Conn, error)) DialOption {
  131. return DialOption{func(do *dialOptions) {
  132. do.dialContext = f
  133. }}
  134. }
  135. // DialDatabase specifies the database to select when dialing a connection.
  136. func DialDatabase(db int) DialOption {
  137. return DialOption{func(do *dialOptions) {
  138. do.db = db
  139. }}
  140. }
  141. // DialPassword specifies the password to use when connecting to
  142. // the Redis server.
  143. func DialPassword(password string) DialOption {
  144. return DialOption{func(do *dialOptions) {
  145. do.password = password
  146. }}
  147. }
  148. // DialUsername specifies the username to use when connecting to
  149. // the Redis server when Redis ACLs are used.
  150. // A DialPassword must also be passed otherwise this option will have no effect.
  151. func DialUsername(username string) DialOption {
  152. return DialOption{func(do *dialOptions) {
  153. do.username = username
  154. }}
  155. }
  156. // DialClientName specifies a client name to be used
  157. // by the Redis server connection.
  158. func DialClientName(name string) DialOption {
  159. return DialOption{func(do *dialOptions) {
  160. do.clientName = name
  161. }}
  162. }
  163. // DialTLSConfig specifies the config to use when a TLS connection is dialed.
  164. // Has no effect when not dialing a TLS connection.
  165. func DialTLSConfig(c *tls.Config) DialOption {
  166. return DialOption{func(do *dialOptions) {
  167. do.tlsConfig = c
  168. }}
  169. }
  170. // DialTLSSkipVerify disables server name verification when connecting over
  171. // TLS. Has no effect when not dialing a TLS connection.
  172. func DialTLSSkipVerify(skip bool) DialOption {
  173. return DialOption{func(do *dialOptions) {
  174. do.skipVerify = skip
  175. }}
  176. }
  177. // DialUseTLS specifies whether TLS should be used when connecting to the
  178. // server. This option is ignore by DialURL.
  179. func DialUseTLS(useTLS bool) DialOption {
  180. return DialOption{func(do *dialOptions) {
  181. do.useTLS = useTLS
  182. }}
  183. }
  184. // Dial connects to the Redis server at the given network and
  185. // address using the specified options.
  186. func Dial(network, address string, options ...DialOption) (Conn, error) {
  187. return DialContext(context.Background(), network, address, options...)
  188. }
  189. type tlsHandshakeTimeoutError struct{}
  190. func (tlsHandshakeTimeoutError) Timeout() bool { return true }
  191. func (tlsHandshakeTimeoutError) Temporary() bool { return true }
  192. func (tlsHandshakeTimeoutError) Error() string { return "TLS handshake timeout" }
  193. // DialContext connects to the Redis server at the given network and
  194. // address using the specified options and context.
  195. func DialContext(ctx context.Context, network, address string, options ...DialOption) (Conn, error) {
  196. do := dialOptions{
  197. dialer: &net.Dialer{
  198. Timeout: time.Second * 30,
  199. KeepAlive: time.Minute * 5,
  200. },
  201. tlsHandshakeTimeout: time.Second * 10,
  202. }
  203. for _, option := range options {
  204. option.f(&do)
  205. }
  206. if do.dialContext == nil {
  207. do.dialContext = do.dialer.DialContext
  208. }
  209. netConn, err := do.dialContext(ctx, network, address)
  210. if err != nil {
  211. return nil, err
  212. }
  213. if do.useTLS {
  214. var tlsConfig *tls.Config
  215. if do.tlsConfig == nil {
  216. tlsConfig = &tls.Config{InsecureSkipVerify: do.skipVerify}
  217. } else {
  218. tlsConfig = cloneTLSConfig(do.tlsConfig)
  219. }
  220. if tlsConfig.ServerName == "" {
  221. host, _, err := net.SplitHostPort(address)
  222. if err != nil {
  223. netConn.Close()
  224. return nil, err
  225. }
  226. tlsConfig.ServerName = host
  227. }
  228. tlsConn := tls.Client(netConn, tlsConfig)
  229. errc := make(chan error, 2) // buffered so we don't block timeout or Handshake
  230. if d := do.tlsHandshakeTimeout; d != 0 {
  231. timer := time.AfterFunc(d, func() {
  232. errc <- tlsHandshakeTimeoutError{}
  233. })
  234. defer timer.Stop()
  235. }
  236. go func() {
  237. errc <- tlsConn.Handshake()
  238. }()
  239. if err := <-errc; err != nil {
  240. // Timeout or Handshake error.
  241. netConn.Close() // nolint: errcheck
  242. return nil, err
  243. }
  244. netConn = tlsConn
  245. }
  246. c := &conn{
  247. conn: netConn,
  248. bw: bufio.NewWriter(netConn),
  249. br: bufio.NewReader(netConn),
  250. readTimeout: do.readTimeout,
  251. writeTimeout: do.writeTimeout,
  252. }
  253. if do.password != "" {
  254. authArgs := make([]interface{}, 0, 2)
  255. if do.username != "" {
  256. authArgs = append(authArgs, do.username)
  257. }
  258. authArgs = append(authArgs, do.password)
  259. if _, err := c.Do("AUTH", authArgs...); err != nil {
  260. netConn.Close()
  261. return nil, err
  262. }
  263. }
  264. if do.clientName != "" {
  265. if _, err := c.Do("CLIENT", "SETNAME", do.clientName); err != nil {
  266. netConn.Close()
  267. return nil, err
  268. }
  269. }
  270. if do.db != 0 {
  271. if _, err := c.Do("SELECT", do.db); err != nil {
  272. netConn.Close()
  273. return nil, err
  274. }
  275. }
  276. return c, nil
  277. }
  278. var pathDBRegexp = regexp.MustCompile(`/(\d*)\z`)
  279. // DialURL connects to a Redis server at the given URL using the Redis
  280. // URI scheme. URLs should follow the draft IANA specification for the
  281. // scheme (https://www.iana.org/assignments/uri-schemes/prov/redis).
  282. func DialURL(rawurl string, options ...DialOption) (Conn, error) {
  283. u, err := url.Parse(rawurl)
  284. if err != nil {
  285. return nil, err
  286. }
  287. if u.Scheme != "redis" && u.Scheme != "rediss" {
  288. return nil, fmt.Errorf("invalid redis URL scheme: %s", u.Scheme)
  289. }
  290. if u.Opaque != "" {
  291. return nil, fmt.Errorf("invalid redis URL, url is opaque: %s", rawurl)
  292. }
  293. // As per the IANA draft spec, the host defaults to localhost and
  294. // the port defaults to 6379.
  295. host, port, err := net.SplitHostPort(u.Host)
  296. if err != nil {
  297. // assume port is missing
  298. host = u.Host
  299. port = "6379"
  300. }
  301. if host == "" {
  302. host = "localhost"
  303. }
  304. address := net.JoinHostPort(host, port)
  305. if u.User != nil {
  306. password, isSet := u.User.Password()
  307. username := u.User.Username()
  308. if isSet {
  309. if username != "" {
  310. // ACL
  311. options = append(options, DialUsername(username), DialPassword(password))
  312. } else {
  313. // requirepass - user-info username:password with blank username
  314. options = append(options, DialPassword(password))
  315. }
  316. } else if username != "" {
  317. // requirepass - redis-cli compatibility which treats as single arg in user-info as a password
  318. options = append(options, DialPassword(username))
  319. }
  320. }
  321. match := pathDBRegexp.FindStringSubmatch(u.Path)
  322. if len(match) == 2 {
  323. db := 0
  324. if len(match[1]) > 0 {
  325. db, err = strconv.Atoi(match[1])
  326. if err != nil {
  327. return nil, fmt.Errorf("invalid database: %s", u.Path[1:])
  328. }
  329. }
  330. if db != 0 {
  331. options = append(options, DialDatabase(db))
  332. }
  333. } else if u.Path != "" {
  334. return nil, fmt.Errorf("invalid database: %s", u.Path[1:])
  335. }
  336. options = append(options, DialUseTLS(u.Scheme == "rediss"))
  337. return Dial("tcp", address, options...)
  338. }
  339. // NewConn returns a new Redigo connection for the given net connection.
  340. func NewConn(netConn net.Conn, readTimeout, writeTimeout time.Duration) Conn {
  341. return &conn{
  342. conn: netConn,
  343. bw: bufio.NewWriter(netConn),
  344. br: bufio.NewReader(netConn),
  345. readTimeout: readTimeout,
  346. writeTimeout: writeTimeout,
  347. }
  348. }
  349. func (c *conn) Close() error {
  350. c.mu.Lock()
  351. err := c.err
  352. if c.err == nil {
  353. c.err = errors.New("redigo: closed")
  354. err = c.conn.Close()
  355. }
  356. c.mu.Unlock()
  357. return err
  358. }
  359. func (c *conn) fatal(err error) error {
  360. c.mu.Lock()
  361. if c.err == nil {
  362. c.err = err
  363. // Close connection to force errors on subsequent calls and to unblock
  364. // other reader or writer.
  365. c.conn.Close()
  366. }
  367. c.mu.Unlock()
  368. return err
  369. }
  370. func (c *conn) Err() error {
  371. c.mu.Lock()
  372. err := c.err
  373. c.mu.Unlock()
  374. return err
  375. }
  376. func (c *conn) writeLen(prefix byte, n int) error {
  377. c.lenScratch[len(c.lenScratch)-1] = '\n'
  378. c.lenScratch[len(c.lenScratch)-2] = '\r'
  379. i := len(c.lenScratch) - 3
  380. for {
  381. c.lenScratch[i] = byte('0' + n%10)
  382. i -= 1
  383. n = n / 10
  384. if n == 0 {
  385. break
  386. }
  387. }
  388. c.lenScratch[i] = prefix
  389. _, err := c.bw.Write(c.lenScratch[i:])
  390. return err
  391. }
  392. func (c *conn) writeString(s string) error {
  393. if err := c.writeLen('$', len(s)); err != nil {
  394. return err
  395. }
  396. if _, err := c.bw.WriteString(s); err != nil {
  397. return err
  398. }
  399. _, err := c.bw.WriteString("\r\n")
  400. return err
  401. }
  402. func (c *conn) writeBytes(p []byte) error {
  403. if err := c.writeLen('$', len(p)); err != nil {
  404. return err
  405. }
  406. if _, err := c.bw.Write(p); err != nil {
  407. return err
  408. }
  409. _, err := c.bw.WriteString("\r\n")
  410. return err
  411. }
  412. func (c *conn) writeInt64(n int64) error {
  413. return c.writeBytes(strconv.AppendInt(c.numScratch[:0], n, 10))
  414. }
  415. func (c *conn) writeFloat64(n float64) error {
  416. return c.writeBytes(strconv.AppendFloat(c.numScratch[:0], n, 'g', -1, 64))
  417. }
  418. func (c *conn) writeCommand(cmd string, args []interface{}) error {
  419. if err := c.writeLen('*', 1+len(args)); err != nil {
  420. return err
  421. }
  422. if err := c.writeString(cmd); err != nil {
  423. return err
  424. }
  425. for _, arg := range args {
  426. if err := c.writeArg(arg, true); err != nil {
  427. return err
  428. }
  429. }
  430. return nil
  431. }
  432. func (c *conn) writeArg(arg interface{}, argumentTypeOK bool) (err error) {
  433. switch arg := arg.(type) {
  434. case string:
  435. return c.writeString(arg)
  436. case []byte:
  437. return c.writeBytes(arg)
  438. case int:
  439. return c.writeInt64(int64(arg))
  440. case int64:
  441. return c.writeInt64(arg)
  442. case float64:
  443. return c.writeFloat64(arg)
  444. case bool:
  445. if arg {
  446. return c.writeString("1")
  447. } else {
  448. return c.writeString("0")
  449. }
  450. case nil:
  451. return c.writeString("")
  452. case Argument:
  453. if argumentTypeOK {
  454. return c.writeArg(arg.RedisArg(), false)
  455. }
  456. // See comment in default clause below.
  457. var buf bytes.Buffer
  458. fmt.Fprint(&buf, arg)
  459. return c.writeBytes(buf.Bytes())
  460. default:
  461. // This default clause is intended to handle builtin numeric types.
  462. // The function should return an error for other types, but this is not
  463. // done for compatibility with previous versions of the package.
  464. var buf bytes.Buffer
  465. fmt.Fprint(&buf, arg)
  466. return c.writeBytes(buf.Bytes())
  467. }
  468. }
  469. type protocolError string
  470. func (pe protocolError) Error() string {
  471. return fmt.Sprintf("redigo: %s (possible server error or unsupported concurrent read by application)", string(pe))
  472. }
  473. // readLine reads a line of input from the RESP stream.
  474. func (c *conn) readLine() ([]byte, error) {
  475. // To avoid allocations, attempt to read the line using ReadSlice. This
  476. // call typically succeeds. The known case where the call fails is when
  477. // reading the output from the MONITOR command.
  478. p, err := c.br.ReadSlice('\n')
  479. if err == bufio.ErrBufferFull {
  480. // The line does not fit in the bufio.Reader's buffer. Fall back to
  481. // allocating a buffer for the line.
  482. buf := append([]byte{}, p...)
  483. for err == bufio.ErrBufferFull {
  484. p, err = c.br.ReadSlice('\n')
  485. buf = append(buf, p...)
  486. }
  487. p = buf
  488. }
  489. if err != nil {
  490. return nil, err
  491. }
  492. i := len(p) - 2
  493. if i < 0 || p[i] != '\r' {
  494. return nil, protocolError("bad response line terminator")
  495. }
  496. return p[:i], nil
  497. }
  498. // parseLen parses bulk string and array lengths.
  499. func parseLen(p []byte) (int, error) {
  500. if len(p) == 0 {
  501. return -1, protocolError("malformed length")
  502. }
  503. if p[0] == '-' && len(p) == 2 && p[1] == '1' {
  504. // handle $-1 and $-1 null replies.
  505. return -1, nil
  506. }
  507. var n int
  508. for _, b := range p {
  509. n *= 10
  510. if b < '0' || b > '9' {
  511. return -1, protocolError("illegal bytes in length")
  512. }
  513. n += int(b - '0')
  514. }
  515. return n, nil
  516. }
  517. // parseInt parses an integer reply.
  518. func parseInt(p []byte) (interface{}, error) {
  519. if len(p) == 0 {
  520. return 0, protocolError("malformed integer")
  521. }
  522. var negate bool
  523. if p[0] == '-' {
  524. negate = true
  525. p = p[1:]
  526. if len(p) == 0 {
  527. return 0, protocolError("malformed integer")
  528. }
  529. }
  530. var n int64
  531. for _, b := range p {
  532. n *= 10
  533. if b < '0' || b > '9' {
  534. return 0, protocolError("illegal bytes in length")
  535. }
  536. n += int64(b - '0')
  537. }
  538. if negate {
  539. n = -n
  540. }
  541. return n, nil
  542. }
  543. var (
  544. okReply interface{} = "OK"
  545. pongReply interface{} = "PONG"
  546. )
  547. func (c *conn) readReply() (interface{}, error) {
  548. line, err := c.readLine()
  549. if err != nil {
  550. return nil, err
  551. }
  552. if len(line) == 0 {
  553. return nil, protocolError("short response line")
  554. }
  555. switch line[0] {
  556. case '+':
  557. switch string(line[1:]) {
  558. case "OK":
  559. // Avoid allocation for frequent "+OK" response.
  560. return okReply, nil
  561. case "PONG":
  562. // Avoid allocation in PING command benchmarks :)
  563. return pongReply, nil
  564. default:
  565. return string(line[1:]), nil
  566. }
  567. case '-':
  568. return Error(line[1:]), nil
  569. case ':':
  570. return parseInt(line[1:])
  571. case '$':
  572. n, err := parseLen(line[1:])
  573. if n < 0 || err != nil {
  574. return nil, err
  575. }
  576. p := make([]byte, n)
  577. _, err = io.ReadFull(c.br, p)
  578. if err != nil {
  579. return nil, err
  580. }
  581. if line, err := c.readLine(); err != nil {
  582. return nil, err
  583. } else if len(line) != 0 {
  584. return nil, protocolError("bad bulk string format")
  585. }
  586. return p, nil
  587. case '*':
  588. n, err := parseLen(line[1:])
  589. if n < 0 || err != nil {
  590. return nil, err
  591. }
  592. r := make([]interface{}, n)
  593. for i := range r {
  594. r[i], err = c.readReply()
  595. if err != nil {
  596. return nil, err
  597. }
  598. }
  599. return r, nil
  600. }
  601. return nil, protocolError("unexpected response line")
  602. }
  603. func (c *conn) Send(cmd string, args ...interface{}) error {
  604. c.mu.Lock()
  605. c.pending += 1
  606. c.mu.Unlock()
  607. if c.writeTimeout != 0 {
  608. if err := c.conn.SetWriteDeadline(time.Now().Add(c.writeTimeout)); err != nil {
  609. return c.fatal(err)
  610. }
  611. }
  612. if err := c.writeCommand(cmd, args); err != nil {
  613. return c.fatal(err)
  614. }
  615. return nil
  616. }
  617. func (c *conn) Flush() error {
  618. if c.writeTimeout != 0 {
  619. if err := c.conn.SetWriteDeadline(time.Now().Add(c.writeTimeout)); err != nil {
  620. return c.fatal(err)
  621. }
  622. }
  623. if err := c.bw.Flush(); err != nil {
  624. return c.fatal(err)
  625. }
  626. return nil
  627. }
  628. func (c *conn) Receive() (interface{}, error) {
  629. return c.ReceiveWithTimeout(c.readTimeout)
  630. }
  631. func (c *conn) ReceiveWithTimeout(timeout time.Duration) (reply interface{}, err error) {
  632. var deadline time.Time
  633. if timeout != 0 {
  634. deadline = time.Now().Add(timeout)
  635. }
  636. if err := c.conn.SetReadDeadline(deadline); err != nil {
  637. return nil, c.fatal(err)
  638. }
  639. if reply, err = c.readReply(); err != nil {
  640. return nil, c.fatal(err)
  641. }
  642. // When using pub/sub, the number of receives can be greater than the
  643. // number of sends. To enable normal use of the connection after
  644. // unsubscribing from all channels, we do not decrement pending to a
  645. // negative value.
  646. //
  647. // The pending field is decremented after the reply is read to handle the
  648. // case where Receive is called before Send.
  649. c.mu.Lock()
  650. if c.pending > 0 {
  651. c.pending -= 1
  652. }
  653. c.mu.Unlock()
  654. if err, ok := reply.(Error); ok {
  655. return nil, err
  656. }
  657. return
  658. }
  659. func (c *conn) Do(cmd string, args ...interface{}) (interface{}, error) {
  660. return c.DoWithTimeout(c.readTimeout, cmd, args...)
  661. }
  662. func (c *conn) DoWithTimeout(readTimeout time.Duration, cmd string, args ...interface{}) (interface{}, error) {
  663. c.mu.Lock()
  664. pending := c.pending
  665. c.pending = 0
  666. c.mu.Unlock()
  667. if cmd == "" && pending == 0 {
  668. return nil, nil
  669. }
  670. if c.writeTimeout != 0 {
  671. if err := c.conn.SetWriteDeadline(time.Now().Add(c.writeTimeout)); err != nil {
  672. return nil, c.fatal(err)
  673. }
  674. }
  675. if cmd != "" {
  676. if err := c.writeCommand(cmd, args); err != nil {
  677. return nil, c.fatal(err)
  678. }
  679. }
  680. if err := c.bw.Flush(); err != nil {
  681. return nil, c.fatal(err)
  682. }
  683. var deadline time.Time
  684. if readTimeout != 0 {
  685. deadline = time.Now().Add(readTimeout)
  686. }
  687. if err := c.conn.SetReadDeadline(deadline); err != nil {
  688. return nil, c.fatal(err)
  689. }
  690. if cmd == "" {
  691. reply := make([]interface{}, pending)
  692. for i := range reply {
  693. r, e := c.readReply()
  694. if e != nil {
  695. return nil, c.fatal(e)
  696. }
  697. reply[i] = r
  698. }
  699. return reply, nil
  700. }
  701. var err error
  702. var reply interface{}
  703. for i := 0; i <= pending; i++ {
  704. var e error
  705. if reply, e = c.readReply(); e != nil {
  706. return nil, c.fatal(e)
  707. }
  708. if e, ok := reply.(Error); ok && err == nil {
  709. err = e
  710. }
  711. }
  712. return reply, err
  713. }