watch.go 21 KB

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