You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
130 lines
2.5 KiB
130 lines
2.5 KiB
package main
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/glycerine/rbuf"
|
|
"github.com/qiniu/log"
|
|
)
|
|
|
|
type BroadcastString struct {
|
|
msgC chan string
|
|
writers map[chan string]bool
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func NewBroadcastString() *BroadcastString {
|
|
b := &BroadcastString{
|
|
msgC: make(chan string, 20), // in case of cmd pipe error
|
|
writers: make(map[chan string]bool, 0),
|
|
}
|
|
go func() {
|
|
for message := range b.msgC {
|
|
b.writeToAll(message)
|
|
}
|
|
}()
|
|
return b
|
|
}
|
|
|
|
func (b *BroadcastString) writeToAll(message string) {
|
|
for c := range b.writers {
|
|
select {
|
|
case c <- message:
|
|
case <-time.After(500 * time.Millisecond):
|
|
log.Println("channel closed, remove from queue")
|
|
delete(b.writers, c)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (b *BroadcastString) WriteMessage(message string) {
|
|
select {
|
|
case b.msgC <- message:
|
|
default:
|
|
}
|
|
}
|
|
|
|
func (b *BroadcastString) Reset() {
|
|
b.msgC = make(chan string, 20)
|
|
}
|
|
|
|
func (b *BroadcastString) AddListener(c chan string) chan string {
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
if c == nil {
|
|
c = make(chan string, 4)
|
|
}
|
|
b.writers[c] = true
|
|
return c
|
|
}
|
|
|
|
func (b *BroadcastString) RemoveListener(c chan string) {
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
delete(b.writers, c)
|
|
}
|
|
|
|
func (b *BroadcastString) Close() {
|
|
close(b.msgC)
|
|
}
|
|
|
|
type BufferBroadcast struct {
|
|
bs *BroadcastString
|
|
|
|
maxSize int
|
|
buf *rbuf.FixedSizeRingBuf // *bytes.Buffer
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func NewBufferBroadcast(size int) *BufferBroadcast {
|
|
if size <= 0 {
|
|
size = 4 * 1024 // 4K
|
|
}
|
|
bufb := &BufferBroadcast{
|
|
maxSize: size,
|
|
bs: NewBroadcastString(),
|
|
buf: rbuf.NewFixedSizeRingBuf(size), // bytes.NewBuffer(nil), // buffer.NewRing(buffer.New(size)),
|
|
}
|
|
bufC := bufb.bs.AddListener(nil)
|
|
go func() {
|
|
for msg := range bufC {
|
|
bufb.buf.Write([]byte(msg))
|
|
}
|
|
}()
|
|
return bufb
|
|
}
|
|
|
|
func (b *BufferBroadcast) Write(data []byte) (n int, err error) {
|
|
b.bs.WriteMessage(string(data)) // should return immediatiely, in case of pipe error
|
|
return len(data), nil
|
|
}
|
|
|
|
func (b *BufferBroadcast) Reset() {
|
|
b.buf.Reset()
|
|
b.bs.Reset()
|
|
}
|
|
|
|
func (b *BufferBroadcast) AddHookFunc(wf func(string) error) chan error {
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
c := b.bs.AddListener(nil)
|
|
errC := make(chan error, 1)
|
|
go func() {
|
|
data := b.buf.Bytes()
|
|
// data, _ := ioutil.ReadAll(b.buf)
|
|
if err := wf(string(data)); err != nil {
|
|
errC <- err
|
|
return
|
|
}
|
|
for msg := range c {
|
|
err := wf(msg)
|
|
if err != nil {
|
|
errC <- err
|
|
break
|
|
}
|
|
}
|
|
}()
|
|
return errC
|
|
}
|