Add Stream messaging struct

This commit is contained in:
Alexandre Iooss 2020-10-17 10:02:38 +02:00
parent 9fc94cf0de
commit 73e6be1274
No known key found for this signature in database
GPG Key ID: 6C79278F3FCDCC02
2 changed files with 107 additions and 0 deletions

77
stream/messaging.go Normal file
View File

@ -0,0 +1,77 @@
// Package stream defines a structure to communication between inputs and outputs
package stream
import "sync"
// Stream makes packages able to subscribe to an incoming stream
type Stream struct {
// Incoming data come from this channel
Broadcast chan<- interface{}
// Use a map to be able to delete an item
outputs map[chan<- interface{}]struct{}
// Mutex to lock this ressource
lock sync.Mutex
}
// New creates a new stream.
func New() *Stream {
s := &Stream{}
broadcast := make(chan interface{}, 64)
s.Broadcast = broadcast
s.outputs = make(map[chan<- interface{}]struct{})
go s.run(broadcast)
return s
}
func (s *Stream) run(broadcast <-chan interface{}) {
for msg := range broadcast {
func() {
s.lock.Lock()
defer s.lock.Unlock()
for output := range s.outputs {
select {
case output <- msg:
default:
// Remove output if failed
delete(s.outputs, output)
close(output)
}
}
}()
}
// Incoming chan has been closed, close all outputs
s.lock.Lock()
defer s.lock.Unlock()
for ch := range s.outputs {
delete(s.outputs, ch)
close(ch)
}
}
// Close the incoming chan, this will also delete all outputs
func (s *Stream) Close() {
close(s.Broadcast)
}
// Register a new output on a stream
func (s *Stream) Register(output chan<- interface{}) {
s.lock.Lock()
defer s.lock.Unlock()
s.outputs[output] = struct{}{}
}
// Unregister removes an output
func (s *Stream) Unregister(output chan<- interface{}) {
s.lock.Lock()
defer s.lock.Unlock()
// Make sure we did not already close this output
_, ok := s.outputs[output]
if ok {
delete(s.outputs, output)
close(output)
}
}

30
stream/messaging_test.go Normal file
View File

@ -0,0 +1,30 @@
package stream
import (
"testing"
)
func TestWithoutOutputs(t *testing.T) {
stream := New()
defer stream.Close()
stream.Broadcast <- "hello world"
}
func TestWithOneOutput(t *testing.T) {
stream := New()
defer stream.Close()
// Register one output
output := make(chan interface{}, 64)
stream.Register(output)
// Try to pass one message
stream.Broadcast <- "hello world"
msg := <-output
if m, ok := msg.(string); !ok || m != "hello world" {
t.Error("Message has wrong type or content")
}
// Unregister
stream.Unregister(output)
}