github.com/grafana/pyroscope@v1.18.0/pkg/metastore/discovery/kuberesolver/stream.go (about)

     1  package kuberesolver
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"io"
     8  	"sync"
     9  
    10  	"google.golang.org/grpc/grpclog"
    11  )
    12  
    13  // Interface can be implemented by anything that knows how to watch and report changes.
    14  type watchInterface interface {
    15  	// Stops watching. Will close the channel returned by ResultChan(). Releases
    16  	// any resources used by the watch.
    17  	Stop()
    18  
    19  	// Returns a chan which will receive all the events. If an error occurs
    20  	// or Stop() is called, this channel will be closed, in which case the
    21  	// watch should be completely cleaned up.
    22  	ResultChan() <-chan Event
    23  }
    24  
    25  // StreamWatcher turns any stream for which you can write a Decoder interface
    26  // into a watch.Interface.
    27  type streamWatcher struct {
    28  	result  chan Event
    29  	r       io.ReadCloser
    30  	decoder *json.Decoder
    31  	sync.Mutex
    32  	stopped bool
    33  }
    34  
    35  // NewStreamWatcher creates a StreamWatcher from the given io.ReadClosers.
    36  func newStreamWatcher(r io.ReadCloser) watchInterface {
    37  	sw := &streamWatcher{
    38  		r:       r,
    39  		decoder: json.NewDecoder(r),
    40  		result:  make(chan Event),
    41  	}
    42  	go sw.receive()
    43  	return sw
    44  }
    45  
    46  // ResultChan implements Interface.
    47  func (sw *streamWatcher) ResultChan() <-chan Event {
    48  	return sw.result
    49  }
    50  
    51  // Stop implements Interface.
    52  func (sw *streamWatcher) Stop() {
    53  	sw.Lock()
    54  	defer sw.Unlock()
    55  	if !sw.stopped {
    56  		sw.stopped = true
    57  		sw.r.Close()
    58  	}
    59  }
    60  
    61  // stopping returns true if Stop() was called previously.
    62  func (sw *streamWatcher) stopping() bool {
    63  	sw.Lock()
    64  	defer sw.Unlock()
    65  	return sw.stopped
    66  }
    67  
    68  // receive reads result from the decoder in a loop and sends down the result channel.
    69  func (sw *streamWatcher) receive() {
    70  	defer close(sw.result)
    71  	defer sw.Stop()
    72  	for {
    73  		obj, err := sw.Decode()
    74  		if err != nil {
    75  			// Ignore expected error.
    76  			if sw.stopping() {
    77  				return
    78  			}
    79  			switch err {
    80  			case io.EOF:
    81  				// watch closed normally
    82  			case context.Canceled:
    83  				// canceled normally
    84  			case io.ErrUnexpectedEOF:
    85  				grpclog.Infof("kuberesolver: Unexpected EOF during watch stream event decoding: %v", err)
    86  			default:
    87  				grpclog.Infof("kuberesolver: Unable to decode an event from the watch stream: %v", err)
    88  			}
    89  			return
    90  		}
    91  		sw.result <- obj
    92  	}
    93  }
    94  
    95  // Decode blocks until it can return the next object in the writer. Returns an error
    96  // if the writer is closed or an object can't be decoded.
    97  func (sw *streamWatcher) Decode() (Event, error) {
    98  	var got Event
    99  	if err := sw.decoder.Decode(&got); err != nil {
   100  		return Event{}, err
   101  	}
   102  	switch got.Type {
   103  	case Added, Modified, Deleted, Error:
   104  		return got, nil
   105  	default:
   106  		return Event{}, fmt.Errorf("got invalid watch event type: %v", got.Type)
   107  	}
   108  }