github.com/gravitational/teleport/api@v0.0.0-20240507183017-3110591cbafc/client/streamwatcher.go (about)

     1  /*
     2  Copyright 2020 Gravitational, Inc.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package client
    18  
    19  import (
    20  	"context"
    21  	"sync"
    22  
    23  	"github.com/gravitational/trace"
    24  
    25  	"github.com/gravitational/teleport/api/client/proto"
    26  	"github.com/gravitational/teleport/api/types"
    27  )
    28  
    29  // NewWatcher returns a new streamWatcher
    30  func (c *Client) NewWatcher(ctx context.Context, watch types.Watch) (types.Watcher, error) {
    31  	cancelCtx, cancel := context.WithCancel(ctx)
    32  	protoWatch := proto.Watch{
    33  		Kinds:               watch.Kinds,
    34  		AllowPartialSuccess: watch.AllowPartialSuccess,
    35  	}
    36  	stream, err := c.grpc.WatchEvents(cancelCtx, &protoWatch)
    37  	if err != nil {
    38  		cancel()
    39  		return nil, trace.Wrap(err)
    40  	}
    41  	w := &streamWatcher{
    42  		stream:  stream,
    43  		ctx:     cancelCtx,
    44  		cancel:  cancel,
    45  		eventsC: make(chan types.Event),
    46  	}
    47  	go w.receiveEvents()
    48  	return w, nil
    49  }
    50  
    51  type streamWatcher struct {
    52  	mu      sync.RWMutex
    53  	stream  proto.AuthService_WatchEventsClient
    54  	ctx     context.Context
    55  	cancel  context.CancelFunc
    56  	eventsC chan types.Event
    57  	err     error
    58  }
    59  
    60  // Error returns the streamWatcher's error
    61  func (w *streamWatcher) Error() error {
    62  	w.mu.RLock()
    63  	defer w.mu.RUnlock()
    64  	if w.err == nil {
    65  		return trace.Wrap(w.ctx.Err())
    66  	}
    67  	return w.err
    68  }
    69  
    70  func (w *streamWatcher) closeWithError(err error) {
    71  	w.mu.Lock()
    72  	defer w.mu.Unlock()
    73  	w.Close()
    74  	w.err = err
    75  }
    76  
    77  // Events returns the streamWatcher's events channel
    78  func (w *streamWatcher) Events() <-chan types.Event {
    79  	return w.eventsC
    80  }
    81  
    82  func (w *streamWatcher) receiveEvents() {
    83  	for {
    84  		event, err := w.stream.Recv()
    85  		if err != nil {
    86  			w.closeWithError(trace.Wrap(err))
    87  			return
    88  		}
    89  		out, err := EventFromGRPC(event)
    90  		if err != nil {
    91  			w.closeWithError(trace.Wrap(err))
    92  			return
    93  		}
    94  		select {
    95  		case w.eventsC <- *out:
    96  		case <-w.Done():
    97  			return
    98  		}
    99  	}
   100  }
   101  
   102  // Done returns a channel that closes once the streamWatcher is Closed
   103  func (w *streamWatcher) Done() <-chan struct{} {
   104  	return w.ctx.Done()
   105  }
   106  
   107  // Close the streamWatcher
   108  func (w *streamWatcher) Close() error {
   109  	w.cancel()
   110  	return nil
   111  }