go-waku/waku/v2/node/broadcast.go

91 lines
2.1 KiB
Go
Raw Normal View History

package node
import (
2021-04-22 00:09:37 +00:00
"github.com/status-im/go-waku/waku/v2/protocol"
)
// Adapted from https://github.com/dustin/go-broadcast/commit/f664265f5a662fb4d1df7f3533b1e8d0e0277120 which was released under MIT license
type broadcaster struct {
2021-04-22 00:09:37 +00:00
input chan *protocol.Envelope
reg chan chan<- *protocol.Envelope
unreg chan chan<- *protocol.Envelope
2021-04-22 00:09:37 +00:00
outputs map[chan<- *protocol.Envelope]bool
}
// The Broadcaster interface describes the main entry points to
// broadcasters.
type Broadcaster interface {
// Register a new channel to receive broadcasts
2021-04-22 00:09:37 +00:00
Register(chan<- *protocol.Envelope)
// Unregister a channel so that it no longer receives broadcasts.
2021-04-22 00:09:37 +00:00
Unregister(chan<- *protocol.Envelope)
// Shut this broadcaster down.
Close() error
// Submit a new object to all subscribers
2021-04-22 00:09:37 +00:00
Submit(*protocol.Envelope)
}
2021-04-22 00:09:37 +00:00
func (b *broadcaster) broadcast(m *protocol.Envelope) {
for ch := range b.outputs {
ch <- m
}
}
func (b *broadcaster) run() {
for {
select {
case m := <-b.input:
b.broadcast(m)
case ch, ok := <-b.reg:
if ok {
b.outputs[ch] = true
} else {
return
}
case ch := <-b.unreg:
delete(b.outputs, ch)
}
}
}
2021-04-22 18:49:52 +00:00
// NewBroadcaster creates a Broadcaster with an specified length
// It's used to register subscriptors that will need to receive
// an Envelope containing a WakuMessage
func NewBroadcaster(buflen int) Broadcaster {
b := &broadcaster{
2021-04-22 00:09:37 +00:00
input: make(chan *protocol.Envelope, buflen),
reg: make(chan chan<- *protocol.Envelope),
unreg: make(chan chan<- *protocol.Envelope),
outputs: make(map[chan<- *protocol.Envelope]bool),
}
go b.run()
return b
}
2021-04-22 18:49:52 +00:00
// Register a subscriptor channel
2021-04-22 00:09:37 +00:00
func (b *broadcaster) Register(newch chan<- *protocol.Envelope) {
b.reg <- newch
}
2021-04-22 18:49:52 +00:00
// Unregister a subscriptor channel
2021-04-22 00:09:37 +00:00
func (b *broadcaster) Unregister(newch chan<- *protocol.Envelope) {
b.unreg <- newch
}
2021-04-22 18:49:52 +00:00
// Closes the broadcaster. Used to stop receiving new subscribers
func (b *broadcaster) Close() error {
close(b.reg)
return nil
}
2021-04-22 18:49:52 +00:00
// Submits an Envelope to be broadcasted among all registered subscriber channels
2021-04-22 00:09:37 +00:00
func (b *broadcaster) Submit(m *protocol.Envelope) {
if b != nil {
b.input <- m
}
}