watch.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807
  1. // Copyright 2016 The etcd Authors
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain 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,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package clientv3
  15. import (
  16. "context"
  17. "fmt"
  18. "sync"
  19. "time"
  20. v3rpc "github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes"
  21. pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
  22. mvccpb "github.com/coreos/etcd/mvcc/mvccpb"
  23. "google.golang.org/grpc"
  24. "google.golang.org/grpc/codes"
  25. "google.golang.org/grpc/metadata"
  26. "google.golang.org/grpc/status"
  27. )
  28. const (
  29. EventTypeDelete = mvccpb.DELETE
  30. EventTypePut = mvccpb.PUT
  31. closeSendErrTimeout = 250 * time.Millisecond
  32. )
  33. type Event mvccpb.Event
  34. type WatchChan <-chan WatchResponse
  35. type Watcher interface {
  36. // Watch watches on a key or prefix. The watched events will be returned
  37. // through the returned channel. If revisions waiting to be sent over the
  38. // watch are compacted, then the watch will be canceled by the server, the
  39. // client will post a compacted error watch response, and the channel will close.
  40. Watch(ctx context.Context, key string, opts ...OpOption) WatchChan
  41. // Close closes the watcher and cancels all watch requests.
  42. Close() error
  43. }
  44. type WatchResponse struct {
  45. Header pb.ResponseHeader
  46. Events []*Event
  47. // CompactRevision is the minimum revision the watcher may receive.
  48. CompactRevision int64
  49. // Canceled is used to indicate watch failure.
  50. // If the watch failed and the stream was about to close, before the channel is closed,
  51. // the channel sends a final response that has Canceled set to true with a non-nil Err().
  52. Canceled bool
  53. // Created is used to indicate the creation of the watcher.
  54. Created bool
  55. closeErr error
  56. // cancelReason is a reason of canceling watch
  57. cancelReason string
  58. }
  59. // IsCreate returns true if the event tells that the key is newly created.
  60. func (e *Event) IsCreate() bool {
  61. return e.Type == EventTypePut && e.Kv.CreateRevision == e.Kv.ModRevision
  62. }
  63. // IsModify returns true if the event tells that a new value is put on existing key.
  64. func (e *Event) IsModify() bool {
  65. return e.Type == EventTypePut && e.Kv.CreateRevision != e.Kv.ModRevision
  66. }
  67. // Err is the error value if this WatchResponse holds an error.
  68. func (wr *WatchResponse) Err() error {
  69. switch {
  70. case wr.closeErr != nil:
  71. return v3rpc.Error(wr.closeErr)
  72. case wr.CompactRevision != 0:
  73. return v3rpc.ErrCompacted
  74. case wr.Canceled:
  75. if len(wr.cancelReason) != 0 {
  76. return v3rpc.Error(status.Error(codes.FailedPrecondition, wr.cancelReason))
  77. }
  78. return v3rpc.ErrFutureRev
  79. }
  80. return nil
  81. }
  82. // IsProgressNotify returns true if the WatchResponse is progress notification.
  83. func (wr *WatchResponse) IsProgressNotify() bool {
  84. return len(wr.Events) == 0 && !wr.Canceled && !wr.Created && wr.CompactRevision == 0 && wr.Header.Revision != 0
  85. }
  86. // watcher implements the Watcher interface
  87. type watcher struct {
  88. remote pb.WatchClient
  89. // mu protects the grpc streams map
  90. mu sync.RWMutex
  91. // streams holds all the active grpc streams keyed by ctx value.
  92. streams map[string]*watchGrpcStream
  93. }
  94. // watchGrpcStream tracks all watch resources attached to a single grpc stream.
  95. type watchGrpcStream struct {
  96. owner *watcher
  97. remote pb.WatchClient
  98. // ctx controls internal remote.Watch requests
  99. ctx context.Context
  100. // ctxKey is the key used when looking up this stream's context
  101. ctxKey string
  102. cancel context.CancelFunc
  103. // substreams holds all active watchers on this grpc stream
  104. substreams map[int64]*watcherStream
  105. // resuming holds all resuming watchers on this grpc stream
  106. resuming []*watcherStream
  107. // reqc sends a watch request from Watch() to the main goroutine
  108. reqc chan *watchRequest
  109. // respc receives data from the watch client
  110. respc chan *pb.WatchResponse
  111. // donec closes to broadcast shutdown
  112. donec chan struct{}
  113. // errc transmits errors from grpc Recv to the watch stream reconnect logic
  114. errc chan error
  115. // closingc gets the watcherStream of closing watchers
  116. closingc chan *watcherStream
  117. // wg is Done when all substream goroutines have exited
  118. wg sync.WaitGroup
  119. // resumec closes to signal that all substreams should begin resuming
  120. resumec chan struct{}
  121. // closeErr is the error that closed the watch stream
  122. closeErr error
  123. }
  124. // watchRequest is issued by the subscriber to start a new watcher
  125. type watchRequest struct {
  126. ctx context.Context
  127. key string
  128. end string
  129. rev int64
  130. // send created notification event if this field is true
  131. createdNotify bool
  132. // progressNotify is for progress updates
  133. progressNotify bool
  134. // filters is the list of events to filter out
  135. filters []pb.WatchCreateRequest_FilterType
  136. // get the previous key-value pair before the event happens
  137. prevKV bool
  138. // retc receives a chan WatchResponse once the watcher is established
  139. retc chan chan WatchResponse
  140. }
  141. // watcherStream represents a registered watcher
  142. type watcherStream struct {
  143. // initReq is the request that initiated this request
  144. initReq watchRequest
  145. // outc publishes watch responses to subscriber
  146. outc chan WatchResponse
  147. // recvc buffers watch responses before publishing
  148. recvc chan *WatchResponse
  149. // donec closes when the watcherStream goroutine stops.
  150. donec chan struct{}
  151. // closing is set to true when stream should be scheduled to shutdown.
  152. closing bool
  153. // id is the registered watch id on the grpc stream
  154. id int64
  155. // buf holds all events received from etcd but not yet consumed by the client
  156. buf []*WatchResponse
  157. }
  158. func NewWatcher(c *Client) Watcher {
  159. return NewWatchFromWatchClient(pb.NewWatchClient(c.conn))
  160. }
  161. func NewWatchFromWatchClient(wc pb.WatchClient) Watcher {
  162. return &watcher{
  163. remote: wc,
  164. streams: make(map[string]*watchGrpcStream),
  165. }
  166. }
  167. // never closes
  168. var valCtxCh = make(chan struct{})
  169. var zeroTime = time.Unix(0, 0)
  170. // ctx with only the values; never Done
  171. type valCtx struct{ context.Context }
  172. func (vc *valCtx) Deadline() (time.Time, bool) { return zeroTime, false }
  173. func (vc *valCtx) Done() <-chan struct{} { return valCtxCh }
  174. func (vc *valCtx) Err() error { return nil }
  175. func (w *watcher) newWatcherGrpcStream(inctx context.Context) *watchGrpcStream {
  176. ctx, cancel := context.WithCancel(&valCtx{inctx})
  177. wgs := &watchGrpcStream{
  178. owner: w,
  179. remote: w.remote,
  180. ctx: ctx,
  181. ctxKey: streamKeyFromCtx(inctx),
  182. cancel: cancel,
  183. substreams: make(map[int64]*watcherStream),
  184. respc: make(chan *pb.WatchResponse),
  185. reqc: make(chan *watchRequest),
  186. donec: make(chan struct{}),
  187. errc: make(chan error, 1),
  188. closingc: make(chan *watcherStream),
  189. resumec: make(chan struct{}),
  190. }
  191. go wgs.run()
  192. return wgs
  193. }
  194. // Watch posts a watch request to run() and waits for a new watcher channel
  195. func (w *watcher) Watch(ctx context.Context, key string, opts ...OpOption) WatchChan {
  196. ow := opWatch(key, opts...)
  197. var filters []pb.WatchCreateRequest_FilterType
  198. if ow.filterPut {
  199. filters = append(filters, pb.WatchCreateRequest_NOPUT)
  200. }
  201. if ow.filterDelete {
  202. filters = append(filters, pb.WatchCreateRequest_NODELETE)
  203. }
  204. wr := &watchRequest{
  205. ctx: ctx,
  206. createdNotify: ow.createdNotify,
  207. key: string(ow.key),
  208. end: string(ow.end),
  209. rev: ow.rev,
  210. progressNotify: ow.progressNotify,
  211. filters: filters,
  212. prevKV: ow.prevKV,
  213. retc: make(chan chan WatchResponse, 1),
  214. }
  215. ok := false
  216. ctxKey := streamKeyFromCtx(ctx)
  217. // find or allocate appropriate grpc watch stream
  218. w.mu.Lock()
  219. if w.streams == nil {
  220. // closed
  221. w.mu.Unlock()
  222. ch := make(chan WatchResponse)
  223. close(ch)
  224. return ch
  225. }
  226. wgs := w.streams[ctxKey]
  227. if wgs == nil {
  228. wgs = w.newWatcherGrpcStream(ctx)
  229. w.streams[ctxKey] = wgs
  230. }
  231. donec := wgs.donec
  232. reqc := wgs.reqc
  233. w.mu.Unlock()
  234. // couldn't create channel; return closed channel
  235. closeCh := make(chan WatchResponse, 1)
  236. // submit request
  237. select {
  238. case reqc <- wr:
  239. ok = true
  240. case <-wr.ctx.Done():
  241. case <-donec:
  242. if wgs.closeErr != nil {
  243. closeCh <- WatchResponse{closeErr: wgs.closeErr}
  244. break
  245. }
  246. // retry; may have dropped stream from no ctxs
  247. return w.Watch(ctx, key, opts...)
  248. }
  249. // receive channel
  250. if ok {
  251. select {
  252. case ret := <-wr.retc:
  253. return ret
  254. case <-ctx.Done():
  255. case <-donec:
  256. if wgs.closeErr != nil {
  257. closeCh <- WatchResponse{closeErr: wgs.closeErr}
  258. break
  259. }
  260. // retry; may have dropped stream from no ctxs
  261. return w.Watch(ctx, key, opts...)
  262. }
  263. }
  264. close(closeCh)
  265. return closeCh
  266. }
  267. func (w *watcher) Close() (err error) {
  268. w.mu.Lock()
  269. streams := w.streams
  270. w.streams = nil
  271. w.mu.Unlock()
  272. for _, wgs := range streams {
  273. if werr := wgs.close(); werr != nil {
  274. err = werr
  275. }
  276. }
  277. return err
  278. }
  279. func (w *watchGrpcStream) close() (err error) {
  280. w.cancel()
  281. <-w.donec
  282. select {
  283. case err = <-w.errc:
  284. default:
  285. }
  286. return toErr(w.ctx, err)
  287. }
  288. func (w *watcher) closeStream(wgs *watchGrpcStream) {
  289. w.mu.Lock()
  290. close(wgs.donec)
  291. wgs.cancel()
  292. if w.streams != nil {
  293. delete(w.streams, wgs.ctxKey)
  294. }
  295. w.mu.Unlock()
  296. }
  297. func (w *watchGrpcStream) addSubstream(resp *pb.WatchResponse, ws *watcherStream) {
  298. if resp.WatchId == -1 {
  299. // failed; no channel
  300. close(ws.recvc)
  301. return
  302. }
  303. ws.id = resp.WatchId
  304. w.substreams[ws.id] = ws
  305. }
  306. func (w *watchGrpcStream) sendCloseSubstream(ws *watcherStream, resp *WatchResponse) {
  307. select {
  308. case ws.outc <- *resp:
  309. case <-ws.initReq.ctx.Done():
  310. case <-time.After(closeSendErrTimeout):
  311. }
  312. close(ws.outc)
  313. }
  314. func (w *watchGrpcStream) closeSubstream(ws *watcherStream) {
  315. // send channel response in case stream was never established
  316. select {
  317. case ws.initReq.retc <- ws.outc:
  318. default:
  319. }
  320. // close subscriber's channel
  321. if closeErr := w.closeErr; closeErr != nil && ws.initReq.ctx.Err() == nil {
  322. go w.sendCloseSubstream(ws, &WatchResponse{closeErr: w.closeErr})
  323. } else if ws.outc != nil {
  324. close(ws.outc)
  325. }
  326. if ws.id != -1 {
  327. delete(w.substreams, ws.id)
  328. return
  329. }
  330. for i := range w.resuming {
  331. if w.resuming[i] == ws {
  332. w.resuming[i] = nil
  333. return
  334. }
  335. }
  336. }
  337. // run is the root of the goroutines for managing a watcher client
  338. func (w *watchGrpcStream) run() {
  339. var wc pb.Watch_WatchClient
  340. var closeErr error
  341. // substreams marked to close but goroutine still running; needed for
  342. // avoiding double-closing recvc on grpc stream teardown
  343. closing := make(map[*watcherStream]struct{})
  344. defer func() {
  345. w.closeErr = closeErr
  346. // shutdown substreams and resuming substreams
  347. for _, ws := range w.substreams {
  348. if _, ok := closing[ws]; !ok {
  349. close(ws.recvc)
  350. closing[ws] = struct{}{}
  351. }
  352. }
  353. for _, ws := range w.resuming {
  354. if _, ok := closing[ws]; ws != nil && !ok {
  355. close(ws.recvc)
  356. closing[ws] = struct{}{}
  357. }
  358. }
  359. w.joinSubstreams()
  360. for range closing {
  361. w.closeSubstream(<-w.closingc)
  362. }
  363. w.wg.Wait()
  364. w.owner.closeStream(w)
  365. }()
  366. // start a stream with the etcd grpc server
  367. if wc, closeErr = w.newWatchClient(); closeErr != nil {
  368. return
  369. }
  370. cancelSet := make(map[int64]struct{})
  371. for {
  372. select {
  373. // Watch() requested
  374. case wreq := <-w.reqc:
  375. outc := make(chan WatchResponse, 1)
  376. ws := &watcherStream{
  377. initReq: *wreq,
  378. id: -1,
  379. outc: outc,
  380. // unbuffered so resumes won't cause repeat events
  381. recvc: make(chan *WatchResponse),
  382. }
  383. ws.donec = make(chan struct{})
  384. w.wg.Add(1)
  385. go w.serveSubstream(ws, w.resumec)
  386. // queue up for watcher creation/resume
  387. w.resuming = append(w.resuming, ws)
  388. if len(w.resuming) == 1 {
  389. // head of resume queue, can register a new watcher
  390. wc.Send(ws.initReq.toPB())
  391. }
  392. // New events from the watch client
  393. case pbresp := <-w.respc:
  394. switch {
  395. case pbresp.Created:
  396. // response to head of queue creation
  397. if ws := w.resuming[0]; ws != nil {
  398. w.addSubstream(pbresp, ws)
  399. w.dispatchEvent(pbresp)
  400. w.resuming[0] = nil
  401. }
  402. if ws := w.nextResume(); ws != nil {
  403. wc.Send(ws.initReq.toPB())
  404. }
  405. case pbresp.Canceled && pbresp.CompactRevision == 0:
  406. delete(cancelSet, pbresp.WatchId)
  407. if ws, ok := w.substreams[pbresp.WatchId]; ok {
  408. // signal to stream goroutine to update closingc
  409. close(ws.recvc)
  410. closing[ws] = struct{}{}
  411. }
  412. default:
  413. // dispatch to appropriate watch stream
  414. if ok := w.dispatchEvent(pbresp); ok {
  415. break
  416. }
  417. // watch response on unexpected watch id; cancel id
  418. if _, ok := cancelSet[pbresp.WatchId]; ok {
  419. break
  420. }
  421. cancelSet[pbresp.WatchId] = struct{}{}
  422. cr := &pb.WatchRequest_CancelRequest{
  423. CancelRequest: &pb.WatchCancelRequest{
  424. WatchId: pbresp.WatchId,
  425. },
  426. }
  427. req := &pb.WatchRequest{RequestUnion: cr}
  428. wc.Send(req)
  429. }
  430. // watch client failed on Recv; spawn another if possible
  431. case err := <-w.errc:
  432. if isHaltErr(w.ctx, err) || toErr(w.ctx, err) == v3rpc.ErrNoLeader {
  433. closeErr = err
  434. return
  435. }
  436. if wc, closeErr = w.newWatchClient(); closeErr != nil {
  437. return
  438. }
  439. if ws := w.nextResume(); ws != nil {
  440. wc.Send(ws.initReq.toPB())
  441. }
  442. cancelSet = make(map[int64]struct{})
  443. case <-w.ctx.Done():
  444. return
  445. case ws := <-w.closingc:
  446. w.closeSubstream(ws)
  447. delete(closing, ws)
  448. if len(w.substreams)+len(w.resuming) == 0 {
  449. // no more watchers on this stream, shutdown
  450. return
  451. }
  452. }
  453. }
  454. }
  455. // nextResume chooses the next resuming to register with the grpc stream. Abandoned
  456. // streams are marked as nil in the queue since the head must wait for its inflight registration.
  457. func (w *watchGrpcStream) nextResume() *watcherStream {
  458. for len(w.resuming) != 0 {
  459. if w.resuming[0] != nil {
  460. return w.resuming[0]
  461. }
  462. w.resuming = w.resuming[1:len(w.resuming)]
  463. }
  464. return nil
  465. }
  466. // dispatchEvent sends a WatchResponse to the appropriate watcher stream
  467. func (w *watchGrpcStream) dispatchEvent(pbresp *pb.WatchResponse) bool {
  468. events := make([]*Event, len(pbresp.Events))
  469. for i, ev := range pbresp.Events {
  470. events[i] = (*Event)(ev)
  471. }
  472. wr := &WatchResponse{
  473. Header: *pbresp.Header,
  474. Events: events,
  475. CompactRevision: pbresp.CompactRevision,
  476. Created: pbresp.Created,
  477. Canceled: pbresp.Canceled,
  478. cancelReason: pbresp.CancelReason,
  479. }
  480. ws, ok := w.substreams[pbresp.WatchId]
  481. if !ok {
  482. return false
  483. }
  484. select {
  485. case ws.recvc <- wr:
  486. case <-ws.donec:
  487. return false
  488. }
  489. return true
  490. }
  491. // serveWatchClient forwards messages from the grpc stream to run()
  492. func (w *watchGrpcStream) serveWatchClient(wc pb.Watch_WatchClient) {
  493. for {
  494. resp, err := wc.Recv()
  495. if err != nil {
  496. select {
  497. case w.errc <- err:
  498. case <-w.donec:
  499. }
  500. return
  501. }
  502. select {
  503. case w.respc <- resp:
  504. case <-w.donec:
  505. return
  506. }
  507. }
  508. }
  509. // serveSubstream forwards watch responses from run() to the subscriber
  510. func (w *watchGrpcStream) serveSubstream(ws *watcherStream, resumec chan struct{}) {
  511. if ws.closing {
  512. panic("created substream goroutine but substream is closing")
  513. }
  514. // nextRev is the minimum expected next revision
  515. nextRev := ws.initReq.rev
  516. resuming := false
  517. defer func() {
  518. if !resuming {
  519. ws.closing = true
  520. }
  521. close(ws.donec)
  522. if !resuming {
  523. w.closingc <- ws
  524. }
  525. w.wg.Done()
  526. }()
  527. emptyWr := &WatchResponse{}
  528. for {
  529. curWr := emptyWr
  530. outc := ws.outc
  531. if len(ws.buf) > 0 {
  532. curWr = ws.buf[0]
  533. } else {
  534. outc = nil
  535. }
  536. select {
  537. case outc <- *curWr:
  538. if ws.buf[0].Err() != nil {
  539. return
  540. }
  541. ws.buf[0] = nil
  542. ws.buf = ws.buf[1:]
  543. case wr, ok := <-ws.recvc:
  544. if !ok {
  545. // shutdown from closeSubstream
  546. return
  547. }
  548. if wr.Created {
  549. if ws.initReq.retc != nil {
  550. ws.initReq.retc <- ws.outc
  551. // to prevent next write from taking the slot in buffered channel
  552. // and posting duplicate create events
  553. ws.initReq.retc = nil
  554. // send first creation event only if requested
  555. if ws.initReq.createdNotify {
  556. ws.outc <- *wr
  557. }
  558. // once the watch channel is returned, a current revision
  559. // watch must resume at the store revision. This is necessary
  560. // for the following case to work as expected:
  561. // wch := m1.Watch("a")
  562. // m2.Put("a", "b")
  563. // <-wch
  564. // If the revision is only bound on the first observed event,
  565. // if wch is disconnected before the Put is issued, then reconnects
  566. // after it is committed, it'll miss the Put.
  567. if ws.initReq.rev == 0 {
  568. nextRev = wr.Header.Revision
  569. }
  570. }
  571. } else {
  572. // current progress of watch; <= store revision
  573. nextRev = wr.Header.Revision
  574. }
  575. if len(wr.Events) > 0 {
  576. nextRev = wr.Events[len(wr.Events)-1].Kv.ModRevision + 1
  577. }
  578. ws.initReq.rev = nextRev
  579. // created event is already sent above,
  580. // watcher should not post duplicate events
  581. if wr.Created {
  582. continue
  583. }
  584. // TODO pause channel if buffer gets too large
  585. ws.buf = append(ws.buf, wr)
  586. case <-w.ctx.Done():
  587. return
  588. case <-ws.initReq.ctx.Done():
  589. return
  590. case <-resumec:
  591. resuming = true
  592. return
  593. }
  594. }
  595. // lazily send cancel message if events on missing id
  596. }
  597. func (w *watchGrpcStream) newWatchClient() (pb.Watch_WatchClient, error) {
  598. // mark all substreams as resuming
  599. close(w.resumec)
  600. w.resumec = make(chan struct{})
  601. w.joinSubstreams()
  602. for _, ws := range w.substreams {
  603. ws.id = -1
  604. w.resuming = append(w.resuming, ws)
  605. }
  606. // strip out nils, if any
  607. var resuming []*watcherStream
  608. for _, ws := range w.resuming {
  609. if ws != nil {
  610. resuming = append(resuming, ws)
  611. }
  612. }
  613. w.resuming = resuming
  614. w.substreams = make(map[int64]*watcherStream)
  615. // connect to grpc stream while accepting watcher cancelation
  616. stopc := make(chan struct{})
  617. donec := w.waitCancelSubstreams(stopc)
  618. wc, err := w.openWatchClient()
  619. close(stopc)
  620. <-donec
  621. // serve all non-closing streams, even if there's a client error
  622. // so that the teardown path can shutdown the streams as expected.
  623. for _, ws := range w.resuming {
  624. if ws.closing {
  625. continue
  626. }
  627. ws.donec = make(chan struct{})
  628. w.wg.Add(1)
  629. go w.serveSubstream(ws, w.resumec)
  630. }
  631. if err != nil {
  632. return nil, v3rpc.Error(err)
  633. }
  634. // receive data from new grpc stream
  635. go w.serveWatchClient(wc)
  636. return wc, nil
  637. }
  638. func (w *watchGrpcStream) waitCancelSubstreams(stopc <-chan struct{}) <-chan struct{} {
  639. var wg sync.WaitGroup
  640. wg.Add(len(w.resuming))
  641. donec := make(chan struct{})
  642. for i := range w.resuming {
  643. go func(ws *watcherStream) {
  644. defer wg.Done()
  645. if ws.closing {
  646. if ws.initReq.ctx.Err() != nil && ws.outc != nil {
  647. close(ws.outc)
  648. ws.outc = nil
  649. }
  650. return
  651. }
  652. select {
  653. case <-ws.initReq.ctx.Done():
  654. // closed ws will be removed from resuming
  655. ws.closing = true
  656. close(ws.outc)
  657. ws.outc = nil
  658. w.wg.Add(1)
  659. go func() {
  660. defer w.wg.Done()
  661. w.closingc <- ws
  662. }()
  663. case <-stopc:
  664. }
  665. }(w.resuming[i])
  666. }
  667. go func() {
  668. defer close(donec)
  669. wg.Wait()
  670. }()
  671. return donec
  672. }
  673. // joinSubstreams waits for all substream goroutines to complete.
  674. func (w *watchGrpcStream) joinSubstreams() {
  675. for _, ws := range w.substreams {
  676. <-ws.donec
  677. }
  678. for _, ws := range w.resuming {
  679. if ws != nil {
  680. <-ws.donec
  681. }
  682. }
  683. }
  684. // openWatchClient retries opening a watch client until success or halt.
  685. // manually retry in case "ws==nil && err==nil"
  686. // TODO: remove FailFast=false
  687. func (w *watchGrpcStream) openWatchClient() (ws pb.Watch_WatchClient, err error) {
  688. for {
  689. select {
  690. case <-w.ctx.Done():
  691. if err == nil {
  692. return nil, w.ctx.Err()
  693. }
  694. return nil, err
  695. default:
  696. }
  697. if ws, err = w.remote.Watch(w.ctx, grpc.FailFast(false)); ws != nil && err == nil {
  698. break
  699. }
  700. if isHaltErr(w.ctx, err) {
  701. return nil, v3rpc.Error(err)
  702. }
  703. }
  704. return ws, nil
  705. }
  706. // toPB converts an internal watch request structure to its protobuf WatchRequest structure.
  707. func (wr *watchRequest) toPB() *pb.WatchRequest {
  708. req := &pb.WatchCreateRequest{
  709. StartRevision: wr.rev,
  710. Key: []byte(wr.key),
  711. RangeEnd: []byte(wr.end),
  712. ProgressNotify: wr.progressNotify,
  713. Filters: wr.filters,
  714. PrevKv: wr.prevKV,
  715. }
  716. cr := &pb.WatchRequest_CreateRequest{CreateRequest: req}
  717. return &pb.WatchRequest{RequestUnion: cr}
  718. }
  719. func streamKeyFromCtx(ctx context.Context) string {
  720. if md, ok := metadata.FromOutgoingContext(ctx); ok {
  721. return fmt.Sprintf("%+v", md)
  722. }
  723. return ""
  724. }