Golang Watch_WatchServer.Context方法代码示例

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 (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()
	}
}