mirror of
https://github.com/lukaszraczylo/traefikoidc.git
synced 2026-10-14 04:07:46 +00:00
508 lines
10 KiB
Go
508 lines
10 KiB
Go
// Basic stream implementation.
|
|
|
|
package miniredis
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"math"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// a Stream is a list of entries, lowest ID (oldest) first, and all "groups".
|
|
type streamKey struct {
|
|
entries []StreamEntry
|
|
groups map[string]*streamGroup
|
|
lastAllocatedID string
|
|
mu sync.Mutex
|
|
}
|
|
|
|
// a StreamEntry is an entry in a stream. The ID is always of the form
|
|
// "123-123".
|
|
// Values is an ordered list of key-value pairs.
|
|
type StreamEntry struct {
|
|
ID string
|
|
Values []string
|
|
}
|
|
|
|
type streamGroup struct {
|
|
stream *streamKey
|
|
lastID string
|
|
pending []pendingEntry
|
|
consumers map[string]*consumer
|
|
}
|
|
|
|
type consumer struct {
|
|
numPendingEntries int
|
|
// these timestamps aren't tracked perfectly
|
|
lastSeen time.Time // "idle" XINFO key
|
|
lastSuccess time.Time // "inactive" XINFO key
|
|
}
|
|
|
|
type pendingEntry struct {
|
|
id string
|
|
consumer string
|
|
deliveryCount int
|
|
lastDelivery time.Time
|
|
}
|
|
|
|
func newStreamKey() *streamKey {
|
|
return &streamKey{
|
|
groups: map[string]*streamGroup{},
|
|
}
|
|
}
|
|
|
|
// generateID doesn't lock the mutex
|
|
func (s *streamKey) generateID(now time.Time) string {
|
|
ts := uint64(now.UnixNano()) / 1_000_000
|
|
|
|
next := fmt.Sprintf("%d-%d", ts, 0)
|
|
if s.lastAllocatedID != "" && streamCmp(s.lastAllocatedID, next) >= 0 {
|
|
last, _ := parseStreamID(s.lastAllocatedID)
|
|
next = fmt.Sprintf("%d-%d", last[0], last[1]+1)
|
|
}
|
|
|
|
lastID := s.lastIDUnlocked()
|
|
if streamCmp(lastID, next) >= 0 {
|
|
last, _ := parseStreamID(lastID)
|
|
next = fmt.Sprintf("%d-%d", last[0], last[1]+1)
|
|
}
|
|
|
|
s.lastAllocatedID = next
|
|
return next
|
|
}
|
|
|
|
// lastID locks the mutex
|
|
func (s *streamKey) lastID() string {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
return s.lastIDUnlocked()
|
|
}
|
|
|
|
// lastID doesn't lock the mutex
|
|
func (s *streamKey) lastIDUnlocked() string {
|
|
if len(s.entries) == 0 {
|
|
return "0-0"
|
|
}
|
|
|
|
return s.entries[len(s.entries)-1].ID
|
|
}
|
|
|
|
func (s *streamKey) copy() *streamKey {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
cpy := &streamKey{
|
|
entries: s.entries,
|
|
}
|
|
groups := map[string]*streamGroup{}
|
|
for k, v := range s.groups {
|
|
gr := v.copy()
|
|
gr.stream = cpy
|
|
groups[k] = gr
|
|
}
|
|
cpy.groups = groups
|
|
return cpy
|
|
}
|
|
|
|
func parseStreamID(id string) ([2]uint64, error) {
|
|
var (
|
|
res [2]uint64
|
|
err error
|
|
)
|
|
parts := strings.SplitN(id, "-", 2)
|
|
res[0], err = strconv.ParseUint(parts[0], 10, 64)
|
|
if err != nil {
|
|
return res, errors.New(msgInvalidStreamID)
|
|
}
|
|
if len(parts) == 2 {
|
|
res[1], err = strconv.ParseUint(parts[1], 10, 64)
|
|
if err != nil {
|
|
return res, errors.New(msgInvalidStreamID)
|
|
}
|
|
}
|
|
return res, nil
|
|
}
|
|
|
|
// compares two stream IDs (of the full format: "123-123"). Returns: -1, 0, 1
|
|
// The given IDs should be valid stream IDs.
|
|
func streamCmp(a, b string) int {
|
|
ap, _ := parseStreamID(a)
|
|
bp, _ := parseStreamID(b)
|
|
|
|
switch {
|
|
case ap[0] < bp[0]:
|
|
return -1
|
|
case ap[0] > bp[0]:
|
|
return 1
|
|
case ap[1] < bp[1]:
|
|
return -1
|
|
case ap[1] > bp[1]:
|
|
return 1
|
|
default:
|
|
return 0
|
|
}
|
|
}
|
|
|
|
// formatStreamID makes a full id ("42-42") out of a partial one ("42")
|
|
func formatStreamID(id string) (string, error) {
|
|
var ts [2]uint64
|
|
parts := strings.SplitN(id, "-", 2)
|
|
|
|
if len(parts) > 0 {
|
|
p, err := strconv.ParseUint(parts[0], 10, 64)
|
|
if err != nil {
|
|
return "", errInvalidEntryID
|
|
}
|
|
ts[0] = p
|
|
}
|
|
if len(parts) > 1 {
|
|
p, err := strconv.ParseUint(parts[1], 10, 64)
|
|
if err != nil {
|
|
return "", errInvalidEntryID
|
|
}
|
|
ts[1] = p
|
|
}
|
|
return fmt.Sprintf("%d-%d", ts[0], ts[1]), nil
|
|
}
|
|
|
|
func formatStreamRangeBound(id string, start bool, reverse bool) (string, error) {
|
|
if id == "-" {
|
|
return "0-0", nil
|
|
}
|
|
|
|
if id == "+" {
|
|
return fmt.Sprintf("%d-%d", uint64(math.MaxUint64), uint64(math.MaxUint64)), nil
|
|
}
|
|
|
|
if id == "0" {
|
|
return "0-0", nil
|
|
}
|
|
|
|
parts := strings.Split(id, "-")
|
|
if len(parts) == 2 {
|
|
return formatStreamID(id)
|
|
}
|
|
|
|
// Incomplete IDs case
|
|
ts, err := strconv.ParseUint(parts[0], 10, 64)
|
|
if err != nil {
|
|
return "", errInvalidEntryID
|
|
}
|
|
|
|
if (!start && !reverse) || (start && reverse) {
|
|
return fmt.Sprintf("%d-%d", ts, uint64(math.MaxUint64)), nil
|
|
}
|
|
|
|
return fmt.Sprintf("%d-%d", ts, 0), nil
|
|
}
|
|
|
|
func reversedStreamEntries(o []StreamEntry) []StreamEntry {
|
|
newStream := make([]StreamEntry, len(o))
|
|
for i, e := range o {
|
|
newStream[len(o)-i-1] = e
|
|
}
|
|
return newStream
|
|
}
|
|
|
|
func (s *streamKey) createGroup(group, id string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
if _, ok := s.groups[group]; ok {
|
|
return errors.New("BUSYGROUP Consumer Group name already exists")
|
|
}
|
|
|
|
if id == "$" {
|
|
id = s.lastIDUnlocked()
|
|
}
|
|
s.groups[group] = &streamGroup{
|
|
stream: s,
|
|
lastID: id,
|
|
consumers: map[string]*consumer{},
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// streamAdd adds an entry to a stream. Returns the new entry ID.
|
|
// If id is empty or "*" the ID will be generated automatically.
|
|
// `values` should have an even length.
|
|
func (s *streamKey) add(entryID string, values []string, now time.Time) (string, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
if entryID == "" || entryID == "*" {
|
|
entryID = s.generateID(now)
|
|
}
|
|
|
|
entryID, err := formatStreamID(entryID)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if entryID == "0-0" {
|
|
return "", errors.New(msgStreamIDZero)
|
|
}
|
|
if streamCmp(s.lastIDUnlocked(), entryID) != -1 {
|
|
return "", errors.New(msgStreamIDTooSmall)
|
|
}
|
|
|
|
s.entries = append(s.entries, StreamEntry{
|
|
ID: entryID,
|
|
Values: values,
|
|
})
|
|
return entryID, nil
|
|
}
|
|
|
|
func (s *streamKey) trim(n int) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
if len(s.entries) > n {
|
|
s.entries = s.entries[len(s.entries)-n:]
|
|
}
|
|
}
|
|
|
|
// trimBefore deletes entries with an id less than the provided id
|
|
// and returns the number of entries deleted
|
|
func (s *streamKey) trimBefore(id string) int {
|
|
s.mu.Lock()
|
|
var delete []string
|
|
for _, entry := range s.entries {
|
|
if streamCmp(entry.ID, id) < 0 {
|
|
delete = append(delete, entry.ID)
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
s.mu.Unlock()
|
|
s.delete(delete)
|
|
return len(delete)
|
|
}
|
|
|
|
// all entries after "id"
|
|
func (s *streamKey) after(id string) []StreamEntry {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
pos := sort.Search(len(s.entries), func(i int) bool {
|
|
return streamCmp(id, s.entries[i].ID) < 0
|
|
})
|
|
return s.entries[pos:]
|
|
}
|
|
|
|
// get a stream entry by ID
|
|
// Also returns the position in the entries slice, if found.
|
|
func (s *streamKey) get(id string) (int, *StreamEntry) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
pos := sort.Search(len(s.entries), func(i int) bool {
|
|
return streamCmp(id, s.entries[i].ID) <= 0
|
|
})
|
|
if len(s.entries) <= pos || s.entries[pos].ID != id {
|
|
return 0, nil
|
|
}
|
|
return pos, &s.entries[pos]
|
|
}
|
|
|
|
func (g *streamGroup) readGroup(
|
|
now time.Time,
|
|
consumerID,
|
|
id string,
|
|
count int,
|
|
noack bool,
|
|
) []StreamEntry {
|
|
if id == ">" {
|
|
// undelivered messages
|
|
msgs := g.stream.after(g.lastID)
|
|
if len(msgs) == 0 {
|
|
return nil
|
|
}
|
|
|
|
if count > 0 && len(msgs) > count {
|
|
msgs = msgs[:count]
|
|
}
|
|
|
|
if !noack {
|
|
shouldAppend := len(g.pending) == 0
|
|
for _, msg := range msgs {
|
|
if !shouldAppend {
|
|
shouldAppend = streamCmp(msg.ID, g.pending[len(g.pending)-1].id) == 1
|
|
}
|
|
|
|
var entry *pendingEntry
|
|
if shouldAppend {
|
|
g.pending = append(g.pending, pendingEntry{})
|
|
entry = &g.pending[len(g.pending)-1]
|
|
} else {
|
|
var pos int
|
|
pos, entry = g.searchPending(msg.ID)
|
|
if entry == nil {
|
|
g.pending = append(g.pending[:pos+1], g.pending[pos:]...)
|
|
entry = &g.pending[pos]
|
|
} else {
|
|
g.consumers[entry.consumer].numPendingEntries--
|
|
}
|
|
}
|
|
|
|
*entry = pendingEntry{
|
|
id: msg.ID,
|
|
consumer: consumerID,
|
|
deliveryCount: 1,
|
|
lastDelivery: now,
|
|
}
|
|
}
|
|
}
|
|
if _, ok := g.consumers[consumerID]; !ok {
|
|
g.consumers[consumerID] = &consumer{}
|
|
}
|
|
g.consumers[consumerID].numPendingEntries += len(msgs)
|
|
g.lastID = msgs[len(msgs)-1].ID
|
|
return msgs
|
|
}
|
|
|
|
// re-deliver messages from the pending list.
|
|
// con := gr.consumers[consumerID]
|
|
msgs := g.pendingAfter(id)
|
|
var res []StreamEntry
|
|
for i, p := range msgs {
|
|
if p.consumer != consumerID {
|
|
continue
|
|
}
|
|
_, entry := g.stream.get(p.id)
|
|
// not found. Weird?
|
|
if entry == nil {
|
|
continue
|
|
}
|
|
p.deliveryCount += 1
|
|
p.lastDelivery = now
|
|
msgs[i] = p
|
|
res = append(res, *entry)
|
|
}
|
|
return res
|
|
}
|
|
|
|
func (g *streamGroup) searchPending(id string) (int, *pendingEntry) {
|
|
pos := sort.Search(len(g.pending), func(i int) bool {
|
|
return streamCmp(id, g.pending[i].id) <= 0
|
|
})
|
|
if pos >= len(g.pending) || g.pending[pos].id != id {
|
|
return pos, nil
|
|
}
|
|
return pos, &g.pending[pos]
|
|
}
|
|
|
|
func (g *streamGroup) ack(ids []string) (int, error) {
|
|
count := 0
|
|
for _, id := range ids {
|
|
if _, err := parseStreamID(id); err != nil {
|
|
return 0, errors.New(msgInvalidStreamID)
|
|
}
|
|
|
|
pos, entry := g.searchPending(id)
|
|
if entry == nil {
|
|
continue
|
|
}
|
|
|
|
consumer := g.consumers[entry.consumer]
|
|
consumer.numPendingEntries--
|
|
|
|
g.pending = append(g.pending[:pos], g.pending[pos+1:]...)
|
|
// don't count deleted entries
|
|
if _, e := g.stream.get(id); e == nil {
|
|
continue
|
|
}
|
|
count++
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
func (s *streamKey) delete(ids []string) (int, error) {
|
|
count := 0
|
|
for _, id := range ids {
|
|
if _, err := parseStreamID(id); err != nil {
|
|
return 0, errors.New(msgInvalidStreamID)
|
|
}
|
|
|
|
i, entry := s.get(id)
|
|
if entry == nil {
|
|
continue
|
|
}
|
|
|
|
s.entries = append(s.entries[:i], s.entries[i+1:]...)
|
|
count++
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
func (g *streamGroup) pendingAfterOrEqual(id string) []pendingEntry {
|
|
pos := sort.Search(len(g.pending), func(i int) bool {
|
|
return streamCmp(id, g.pending[i].id) <= 0
|
|
})
|
|
return g.pending[pos:]
|
|
}
|
|
|
|
func (g *streamGroup) pendingAfter(id string) []pendingEntry {
|
|
pos := sort.Search(len(g.pending), func(i int) bool {
|
|
return streamCmp(id, g.pending[i].id) < 0
|
|
})
|
|
return g.pending[pos:]
|
|
}
|
|
|
|
func (g *streamGroup) pendingCount(consumer string) int {
|
|
n := 0
|
|
for _, p := range g.activePending() {
|
|
if p.consumer == consumer {
|
|
n++
|
|
}
|
|
}
|
|
return n
|
|
}
|
|
|
|
// pending entries without the entries deleted from the group
|
|
func (g *streamGroup) activePending() []pendingEntry {
|
|
var pe []pendingEntry
|
|
for _, p := range g.pending {
|
|
// drop deleted ones
|
|
if _, e := g.stream.get(p.id); e == nil {
|
|
continue
|
|
}
|
|
p := p
|
|
pe = append(pe, p)
|
|
}
|
|
return pe
|
|
}
|
|
|
|
func (g *streamGroup) copy() *streamGroup {
|
|
cns := map[string]*consumer{}
|
|
for k, v := range g.consumers {
|
|
c := *v
|
|
cns[k] = &c
|
|
}
|
|
return &streamGroup{
|
|
// don't copy stream
|
|
lastID: g.lastID,
|
|
pending: g.pending,
|
|
consumers: cns,
|
|
}
|
|
}
|
|
|
|
func (g *streamGroup) setLastSeen(c string, t time.Time) {
|
|
cons, ok := g.consumers[c]
|
|
if !ok {
|
|
cons = &consumer{}
|
|
}
|
|
cons.lastSeen = t
|
|
g.consumers[c] = cons
|
|
}
|
|
|
|
func (g *streamGroup) setLastSuccess(c string, t time.Time) {
|
|
g.setLastSeen(c, t)
|
|
g.consumers[c].lastSuccess = t
|
|
}
|