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

     1  package stream
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/quickfeed/quickfeed/qf"
     7  )
     8  
     9  // StreamServices contain all available stream services.
    10  // Each service is unique to a specific type.
    11  // The services may be used to send data to connected clients.
    12  // To add a new service, add a new field to this struct and
    13  // initialize the service in the NewStreamServices function.
    14  type StreamServices struct {
    15  	Submission *Service[uint64, qf.Submission]
    16  }
    17  
    18  // NewStreamServices creates a new StreamServices.
    19  func NewStreamServices() *StreamServices {
    20  	return &StreamServices{
    21  		Submission: NewService[uint64, qf.Submission](),
    22  	}
    23  }
    24  
    25  // ID is the allowed type for stream IDs.
    26  type ID interface {
    27  	uint64 | string
    28  }
    29  
    30  // Service[K ID, V any] is a type specific stream service.
    31  // It also contains a map of streams that are currently connected.
    32  type Service[K ID, V any] struct {
    33  	mu sync.Mutex
    34  	// The map of streams.
    35  	streams map[K]StreamInterface[V]
    36  }
    37  
    38  // NewService creates a new service.
    39  func NewService[K ID, V any]() *Service[K, V] {
    40  	return &Service[K, V]{
    41  		streams: make(map[K]StreamInterface[V]),
    42  	}
    43  }
    44  
    45  // SendTo sends data to connected clients with the given IDs.
    46  // If no ID is given, data is sent to all connected clients.
    47  // Unconnected clients are ignored and will not receive the data.
    48  func (s *Service[K, V]) SendTo(data *V, ids ...K) {
    49  	s.mu.Lock()
    50  	defer s.mu.Unlock()
    51  	if len(ids) == 0 {
    52  		// Broadcast to all clients.
    53  		for _, stream := range s.streams {
    54  			stream.Send(data)
    55  		}
    56  		return
    57  	}
    58  
    59  	for _, id := range ids {
    60  		stream, ok := s.streams[id]
    61  		if !ok {
    62  			continue
    63  		}
    64  		stream.Send(data)
    65  	}
    66  }
    67  
    68  // Add adds a new stream for the given identifier.
    69  // The identifier may be a user ID or an external application ID.
    70  func (s *Service[K, V]) Add(stream StreamInterface[V], id K) {
    71  	s.mu.Lock()
    72  	defer s.mu.Unlock()
    73  	// Delete the stream if it already exists.
    74  	s.internalRemove(id)
    75  	// Add the stream to the map.
    76  	s.streams[id] = stream
    77  }
    78  
    79  // internalRemove removes a stream from the service.
    80  // This closes the stream and removes it from the map.
    81  // This function must only be called when holding the mutex.
    82  func (s *Service[K, V]) internalRemove(id K) {
    83  	if stream, ok := s.streams[id]; ok {
    84  		stream.Close()
    85  		delete(s.streams, id)
    86  	}
    87  }
    88  
    89  // Close closes all streams in the service.
    90  func (s *Service[K, V]) Close() {
    91  	s.mu.Lock()
    92  	defer s.mu.Unlock()
    93  	for id := range s.streams {
    94  		s.internalRemove(id)
    95  	}
    96  }
    97  
    98  // CloseBy closes a stream for the given ID, if any exists.
    99  func (s *Service[K, V]) CloseBy(id K) {
   100  	s.mu.Lock()
   101  	defer s.mu.Unlock()
   102  	s.internalRemove(id)
   103  }