github.com/quickfeed/quickfeed@v0.0.0-20240507093252-ed8ca812a09c/web/stream/stream.go (about)

     1  package stream
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  
     8  	"connectrpc.com/connect"
     9  )
    10  
    11  type StreamInterface[T any] interface {
    12  	Send(data *T)
    13  	Run() error
    14  	Close()
    15  }
    16  
    17  // Stream wraps a connect.ServerStream.
    18  type Stream[T any] struct {
    19  	mu sync.Mutex
    20  	// stream is the underlying connect stream
    21  	// that does the actual transfer of data
    22  	// between the server and a client
    23  	stream *connect.ServerStream[T]
    24  	// context is the context of the stream
    25  	ctx context.Context
    26  	// The channel that we listen to for any
    27  	// new data that we need to send to the client.
    28  	ch chan *T
    29  	// closed is a flag that indicates whether
    30  	// the stream has been closed.
    31  	closed bool
    32  }
    33  
    34  // newStream creates a new stream.
    35  func NewStream[T any](ctx context.Context, st *connect.ServerStream[T]) *Stream[T] {
    36  	return &Stream[T]{
    37  		stream: st,
    38  		ctx:    ctx,
    39  		ch:     make(chan *T),
    40  	}
    41  }
    42  
    43  // Close closes the stream.
    44  func (s *Stream[T]) Close() {
    45  	s.mu.Lock()
    46  	defer s.mu.Unlock()
    47  	if !s.closed {
    48  		close(s.ch)
    49  	}
    50  	s.closed = true
    51  }
    52  
    53  // Run runs the stream.
    54  // Run will block until the stream is closed.
    55  func (s *Stream[T]) Run() error {
    56  	defer s.Close()
    57  	for {
    58  		select {
    59  		case <-s.ctx.Done():
    60  			return s.ctx.Err()
    61  		case data, ok := <-s.ch:
    62  			if !ok {
    63  				return connect.NewError(connect.CodeCanceled, fmt.Errorf("stream closed"))
    64  			}
    65  			if err := s.stream.Send(data); err != nil {
    66  				return err
    67  			}
    68  		}
    69  	}
    70  }
    71  
    72  // Send sends data to this stream's connected client.
    73  func (s *Stream[T]) Send(data *T) {
    74  	s.mu.Lock()
    75  	defer s.mu.Unlock()
    76  	if !s.closed {
    77  		s.ch <- data
    78  	}
    79  }