[98] | 1 | package soju
|
---|
[50] | 2 |
|
---|
| 3 | import (
|
---|
[138] | 4 | "fmt"
|
---|
[57] | 5 | "sync"
|
---|
| 6 |
|
---|
[50] | 7 | "gopkg.in/irc.v3"
|
---|
| 8 | )
|
---|
| 9 |
|
---|
| 10 | // Ring implements a single producer, multiple consumer ring buffer. The ring
|
---|
| 11 | // buffer size is fixed. The ring buffer is stored in memory.
|
---|
| 12 | type Ring struct {
|
---|
[57] | 13 | buffer []*irc.Message
|
---|
| 14 | cap uint64
|
---|
[50] | 15 |
|
---|
[57] | 16 | lock sync.Mutex
|
---|
| 17 | cur uint64
|
---|
| 18 | consumers []*RingConsumer
|
---|
[50] | 19 | }
|
---|
| 20 |
|
---|
[59] | 21 | // NewRing creates a new ring buffer.
|
---|
[50] | 22 | func NewRing(capacity int) *Ring {
|
---|
| 23 | return &Ring{
|
---|
[57] | 24 | buffer: make([]*irc.Message, capacity),
|
---|
| 25 | cap: uint64(capacity),
|
---|
[50] | 26 | }
|
---|
| 27 | }
|
---|
| 28 |
|
---|
[59] | 29 | // Produce appends a new message to the ring buffer.
|
---|
[50] | 30 | func (r *Ring) Produce(msg *irc.Message) {
|
---|
[57] | 31 | r.lock.Lock()
|
---|
| 32 | defer r.lock.Unlock()
|
---|
| 33 |
|
---|
[50] | 34 | i := int(r.cur % r.cap)
|
---|
| 35 | r.buffer[i] = msg
|
---|
| 36 | r.cur++
|
---|
| 37 |
|
---|
[57] | 38 | for _, consumer := range r.consumers {
|
---|
| 39 | select {
|
---|
| 40 | case consumer.ch <- struct{}{}:
|
---|
| 41 | // This space is intentionally left blank
|
---|
| 42 | default:
|
---|
| 43 | // The channel already has a pending item
|
---|
| 44 | }
|
---|
[51] | 45 | }
|
---|
[57] | 46 | }
|
---|
[51] | 47 |
|
---|
[59] | 48 | // NewConsumer creates a new ring buffer consumer.
|
---|
| 49 | //
|
---|
| 50 | // If seq is nil, the consumer will get messages starting from the last
|
---|
| 51 | // producer message. If seq is non-nil, the consumer will get messages starting
|
---|
| 52 | // from the specified history sequence number (see RingConsumer.Close).
|
---|
| 53 | //
|
---|
| 54 | // The returned channel yields a value each time the consumer has a new message
|
---|
| 55 | // available. Consume should be called to drain the consumer.
|
---|
| 56 | //
|
---|
| 57 | // The consumer can only be used from a single goroutine.
|
---|
| 58 | func (r *Ring) NewConsumer(seq *uint64) (*RingConsumer, <-chan struct{}) {
|
---|
[57] | 59 | consumer := &RingConsumer{
|
---|
[50] | 60 | ring: r,
|
---|
[57] | 61 | ch: make(chan struct{}, 1),
|
---|
[50] | 62 | }
|
---|
[57] | 63 |
|
---|
| 64 | r.lock.Lock()
|
---|
| 65 | if seq != nil {
|
---|
| 66 | consumer.cur = *seq
|
---|
| 67 | } else {
|
---|
| 68 | consumer.cur = r.cur
|
---|
| 69 | }
|
---|
| 70 | if consumer.diff() > 0 {
|
---|
| 71 | consumer.ch <- struct{}{}
|
---|
| 72 | }
|
---|
| 73 | r.consumers = append(r.consumers, consumer)
|
---|
| 74 | r.lock.Unlock()
|
---|
| 75 |
|
---|
| 76 | return consumer, consumer.ch
|
---|
[50] | 77 | }
|
---|
| 78 |
|
---|
[59] | 79 | // RingConsumer is a ring buffer consumer.
|
---|
[50] | 80 | type RingConsumer struct {
|
---|
[57] | 81 | ring *Ring
|
---|
| 82 | cur uint64
|
---|
| 83 | ch chan struct{}
|
---|
| 84 | closed bool
|
---|
[50] | 85 | }
|
---|
| 86 |
|
---|
[57] | 87 | // diff returns the number of pending messages. It assumes the Ring is locked.
|
---|
| 88 | func (rc *RingConsumer) diff() uint64 {
|
---|
[50] | 89 | if rc.cur > rc.ring.cur {
|
---|
[138] | 90 | panic(fmt.Sprintf("soju: consumer cursor (%v) greater than producer cursor (%v)", rc.cur, rc.ring.cur))
|
---|
[50] | 91 | }
|
---|
| 92 | return rc.ring.cur - rc.cur
|
---|
| 93 | }
|
---|
| 94 |
|
---|
[59] | 95 | // Peek returns the next pending message if any without consuming it. A nil
|
---|
| 96 | // message is returned if no message is available.
|
---|
[50] | 97 | func (rc *RingConsumer) Peek() *irc.Message {
|
---|
[57] | 98 | if rc.closed {
|
---|
[98] | 99 | panic("soju: RingConsumer.Peek called after Close")
|
---|
[57] | 100 | }
|
---|
| 101 |
|
---|
| 102 | rc.ring.lock.Lock()
|
---|
| 103 | defer rc.ring.lock.Unlock()
|
---|
| 104 |
|
---|
| 105 | diff := rc.diff()
|
---|
[50] | 106 | if diff == 0 {
|
---|
| 107 | return nil
|
---|
| 108 | }
|
---|
| 109 | if diff > rc.ring.cap {
|
---|
| 110 | // Consumer drops diff - cap entries
|
---|
| 111 | rc.cur = rc.ring.cur - rc.ring.cap
|
---|
| 112 | }
|
---|
| 113 | i := int(rc.cur % rc.ring.cap)
|
---|
| 114 | msg := rc.ring.buffer[i]
|
---|
| 115 | if msg == nil {
|
---|
[138] | 116 | panic(fmt.Sprintf("soju: unexpected nil ring buffer entry at index %v", i))
|
---|
[50] | 117 | }
|
---|
| 118 | return msg
|
---|
| 119 | }
|
---|
| 120 |
|
---|
[59] | 121 | // Consume consumes and returns the next pending message. A nil message is
|
---|
| 122 | // returned if no message is available.
|
---|
[50] | 123 | func (rc *RingConsumer) Consume() *irc.Message {
|
---|
| 124 | msg := rc.Peek()
|
---|
| 125 | if msg != nil {
|
---|
| 126 | rc.cur++
|
---|
| 127 | }
|
---|
| 128 | return msg
|
---|
| 129 | }
|
---|
[51] | 130 |
|
---|
[59] | 131 | // Close stops consuming messages. The consumer channel will be closed. The
|
---|
| 132 | // current history sequence number is returned. It can be provided later as an
|
---|
| 133 | // argument to Ring.NewConsumer to resume the message stream.
|
---|
[57] | 134 | func (rc *RingConsumer) Close() uint64 {
|
---|
| 135 | rc.ring.lock.Lock()
|
---|
| 136 | for i := range rc.ring.consumers {
|
---|
| 137 | if rc.ring.consumers[i] == rc {
|
---|
| 138 | rc.ring.consumers = append(rc.ring.consumers[:i], rc.ring.consumers[i+1:]...)
|
---|
| 139 | break
|
---|
| 140 | }
|
---|
| 141 | }
|
---|
| 142 | rc.ring.lock.Unlock()
|
---|
| 143 |
|
---|
| 144 | close(rc.ch)
|
---|
| 145 | rc.closed = true
|
---|
| 146 | return rc.cur
|
---|
[51] | 147 | }
|
---|