[50] | 1 | package jounce
|
---|
| 2 |
|
---|
| 3 | import (
|
---|
[57] | 4 | "sync"
|
---|
| 5 |
|
---|
[50] | 6 | "gopkg.in/irc.v3"
|
---|
| 7 | )
|
---|
| 8 |
|
---|
| 9 | // Ring implements a single producer, multiple consumer ring buffer. The ring
|
---|
| 10 | // buffer size is fixed. The ring buffer is stored in memory.
|
---|
| 11 | type Ring struct {
|
---|
[57] | 12 | buffer []*irc.Message
|
---|
| 13 | cap uint64
|
---|
[50] | 14 |
|
---|
[57] | 15 | lock sync.Mutex
|
---|
| 16 | cur uint64
|
---|
| 17 | consumers []*RingConsumer
|
---|
[50] | 18 | }
|
---|
| 19 |
|
---|
| 20 | func NewRing(capacity int) *Ring {
|
---|
| 21 | return &Ring{
|
---|
[57] | 22 | buffer: make([]*irc.Message, capacity),
|
---|
| 23 | cap: uint64(capacity),
|
---|
[50] | 24 | }
|
---|
| 25 | }
|
---|
| 26 |
|
---|
| 27 | func (r *Ring) Produce(msg *irc.Message) {
|
---|
[57] | 28 | r.lock.Lock()
|
---|
| 29 | defer r.lock.Unlock()
|
---|
| 30 |
|
---|
[50] | 31 | i := int(r.cur % r.cap)
|
---|
| 32 | r.buffer[i] = msg
|
---|
| 33 | r.cur++
|
---|
| 34 |
|
---|
[57] | 35 | for _, consumer := range r.consumers {
|
---|
| 36 | select {
|
---|
| 37 | case consumer.ch <- struct{}{}:
|
---|
| 38 | // This space is intentionally left blank
|
---|
| 39 | default:
|
---|
| 40 | // The channel already has a pending item
|
---|
| 41 | }
|
---|
[51] | 42 | }
|
---|
[57] | 43 | }
|
---|
[51] | 44 |
|
---|
[57] | 45 | func (r *Ring) Consumer(seq *uint64) (*RingConsumer, <-chan struct{}) {
|
---|
| 46 | consumer := &RingConsumer{
|
---|
[50] | 47 | ring: r,
|
---|
[57] | 48 | ch: make(chan struct{}, 1),
|
---|
[50] | 49 | }
|
---|
[57] | 50 |
|
---|
| 51 | r.lock.Lock()
|
---|
| 52 | if seq != nil {
|
---|
| 53 | consumer.cur = *seq
|
---|
| 54 | } else {
|
---|
| 55 | consumer.cur = r.cur
|
---|
| 56 | }
|
---|
| 57 | if consumer.diff() > 0 {
|
---|
| 58 | consumer.ch <- struct{}{}
|
---|
| 59 | }
|
---|
| 60 | r.consumers = append(r.consumers, consumer)
|
---|
| 61 | r.lock.Unlock()
|
---|
| 62 |
|
---|
| 63 | return consumer, consumer.ch
|
---|
[50] | 64 | }
|
---|
| 65 |
|
---|
| 66 | type RingConsumer struct {
|
---|
[57] | 67 | ring *Ring
|
---|
| 68 | cur uint64
|
---|
| 69 | ch chan struct{}
|
---|
| 70 | closed bool
|
---|
[50] | 71 | }
|
---|
| 72 |
|
---|
[57] | 73 | // diff returns the number of pending messages. It assumes the Ring is locked.
|
---|
| 74 | func (rc *RingConsumer) diff() uint64 {
|
---|
[50] | 75 | if rc.cur > rc.ring.cur {
|
---|
| 76 | panic("jounce: consumer cursor greater than producer cursor")
|
---|
| 77 | }
|
---|
| 78 | return rc.ring.cur - rc.cur
|
---|
| 79 | }
|
---|
| 80 |
|
---|
| 81 | func (rc *RingConsumer) Peek() *irc.Message {
|
---|
[57] | 82 | if rc.closed {
|
---|
| 83 | panic("jounce: RingConsumer.Peek called after Close")
|
---|
| 84 | }
|
---|
| 85 |
|
---|
| 86 | rc.ring.lock.Lock()
|
---|
| 87 | defer rc.ring.lock.Unlock()
|
---|
| 88 |
|
---|
| 89 | diff := rc.diff()
|
---|
[50] | 90 | if diff == 0 {
|
---|
| 91 | return nil
|
---|
| 92 | }
|
---|
| 93 | if diff > rc.ring.cap {
|
---|
| 94 | // Consumer drops diff - cap entries
|
---|
| 95 | rc.cur = rc.ring.cur - rc.ring.cap
|
---|
| 96 | }
|
---|
| 97 | i := int(rc.cur % rc.ring.cap)
|
---|
| 98 | msg := rc.ring.buffer[i]
|
---|
| 99 | if msg == nil {
|
---|
| 100 | panic("jounce: unexpected nil ring buffer entry")
|
---|
| 101 | }
|
---|
| 102 | return msg
|
---|
| 103 | }
|
---|
| 104 |
|
---|
| 105 | func (rc *RingConsumer) Consume() *irc.Message {
|
---|
| 106 | msg := rc.Peek()
|
---|
| 107 | if msg != nil {
|
---|
| 108 | rc.cur++
|
---|
| 109 | }
|
---|
| 110 | return msg
|
---|
| 111 | }
|
---|
[51] | 112 |
|
---|
[57] | 113 | func (rc *RingConsumer) Close() uint64 {
|
---|
| 114 | rc.ring.lock.Lock()
|
---|
| 115 | for i := range rc.ring.consumers {
|
---|
| 116 | if rc.ring.consumers[i] == rc {
|
---|
| 117 | rc.ring.consumers = append(rc.ring.consumers[:i], rc.ring.consumers[i+1:]...)
|
---|
| 118 | break
|
---|
| 119 | }
|
---|
| 120 | }
|
---|
| 121 | rc.ring.lock.Unlock()
|
---|
| 122 |
|
---|
| 123 | close(rc.ch)
|
---|
| 124 | rc.closed = true
|
---|
| 125 | return rc.cur
|
---|
[51] | 126 | }
|
---|