func sendLoop(stream pb.Watch_WatchServer, watcher storage.Watcher, closec chan struct{}) {
for {
select {
case e, ok := <-watcher.Chan():
if !ok {
return
}
err := stream.Send(&pb.WatchResponse{Event: &e})
storage.ReportEventReceived()
if err != nil {
return
}
case <-closec:
// drain the chan to clean up pending events
for {
_, ok := <-watcher.Chan()
if !ok {
return
}
storage.ReportEventReceived()
}
return
}
}
}
func (ws *watchServer) Watch(stream pb.Watch_WatchServer) error {
closec := make(chan struct{})
defer close(closec)
watcher := ws.watchable.NewWatcher()
defer watcher.Close()
go sendLoop(stream, watcher, closec)
for {
req, err := stream.Recv()
if err == io.EOF {
return nil
}
if err != nil {
return err
}
var prefix bool
toWatch := req.Key
if len(req.Key) == 0 {
toWatch = req.Prefix
prefix = true
}
// TODO: support cancellation
watcher.Watch(toWatch, prefix, req.StartRevision)
}
}
func (ws *watchServer) Watch(stream pb.Watch_WatchServer) error {
sws := serverWatchStream{
clusterID: ws.clusterID,
memberID: ws.memberID,
raftTimer: ws.raftTimer,
gRPCStream: stream,
watchStream: ws.watchable.NewWatchStream(),
// chan for sending control response like watcher created and canceled.
ctrlStream: make(chan *pb.WatchResponse, ctrlStreamBufLen),
progress: make(map[mvcc.WatchID]bool),
closec: make(chan struct{}),
}
go sws.sendLoop()
errc := make(chan error, 1)
go func() {
errc <- sws.recvLoop()
sws.close()
}()
select {
case err := <-errc:
return err
case <-stream.Context().Done():
err := stream.Context().Err()
// the only server-side cancellation is noleader for now.
if err == context.Canceled {
return rpctypes.ErrGRPCNoLeader
}
return err
}
}
func sendLoop(stream pb.Watch_WatchServer, watcher storage.Watcher, closec chan struct{}) {
for {
select {
case e := <-watcher.Chan():
err := stream.Send(&pb.WatchResponse{Event: &e})
if err != nil {
return
}
case <-closec:
return
}
}
}
func (ws *watchServer) Watch(stream pb.Watch_WatchServer) (err error) {
sws := serverWatchStream{
clusterID: ws.clusterID,
memberID: ws.memberID,
raftTimer: ws.raftTimer,
watchable: ws.watchable,
gRPCStream: stream,
watchStream: ws.watchable.NewWatchStream(),
// chan for sending control response like watcher created and canceled.
ctrlStream: make(chan *pb.WatchResponse, ctrlStreamBufLen),
progress: make(map[mvcc.WatchID]bool),
prevKV: make(map[mvcc.WatchID]bool),
closec: make(chan struct{}),
}
sws.wg.Add(1)
go func() {
sws.sendLoop()
sws.wg.Done()
}()
errc := make(chan error, 1)
// Ideally recvLoop would also use sws.wg to signal its completion
// but when stream.Context().Done() is closed, the stream's recv
// may continue to block since it uses a different context, leading to
// deadlock when calling sws.close().
go func() {
if rerr := sws.recvLoop(); rerr != nil {
errc <- rerr
}
}()
select {
case err = <-errc:
close(sws.ctrlStream)
case <-stream.Context().Done():
err = stream.Context().Err()
// the only server-side cancellation is noleader for now.
if err == context.Canceled {
err = rpctypes.ErrGRPCNoLeader
}
}
sws.close()
return err
}
func (wp *watchProxy) Watch(stream pb.Watch_WatchServer) (err error) {
wp.mu.Lock()
select {
case <-wp.ctx.Done():
wp.mu.Unlock()
return
default:
wp.wg.Add(1)
}
wp.mu.Unlock()
ctx, cancel := context.WithCancel(stream.Context())
wps := &watchProxyStream{
ranges: wp.ranges,
watchers: make(map[int64]*watcher),
stream: stream,
watchCh: make(chan *pb.WatchResponse, 1024),
ctx: ctx,
cancel: cancel,
}
var leaderc <-chan struct{}
if md, ok := metadata.FromContext(stream.Context()); ok {
v := md[rpctypes.MetadataRequireLeaderKey]
if len(v) > 0 && v[0] == rpctypes.MetadataHasLeader {
leaderc = wp.lostLeaderNotify()
}
}
// post to stopc => terminate server stream; can't use a waitgroup
// since all goroutines will only terminate after Watch() exits.
stopc := make(chan struct{}, 3)
go func() {
defer func() { stopc <- struct{}{} }()
wps.recvLoop()
}()
go func() {
defer func() { stopc <- struct{}{} }()
wps.sendLoop()
}()
if leaderc != nil {
go func() {
defer func() { stopc <- struct{}{} }()
select {
case <-leaderc:
case <-ctx.Done():
}
}()
}
<-stopc
// recv/send may only shutdown after function exits;
// goroutine notifies proxy that stream is through
go func() {
if leaderc != nil {
<-stopc
}
<-stopc
wps.close()
wp.wg.Done()
}()
select {
case <-leaderc:
return rpctypes.ErrNoLeader
default:
return wps.ctx.Err()
}
}