90 lines
2.7 KiB
Go
90 lines
2.7 KiB
Go
package naffka
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
)
|
|
|
|
// A MemoryDatabase stores the message history as arrays in memory.
|
|
// It can be used to run unit tests.
|
|
// If the process is stopped then any messages that haven't been
|
|
// processed by a consumer are lost forever and all offsets become
|
|
// invalid.
|
|
type MemoryDatabase struct {
|
|
topicsMutex sync.Mutex
|
|
topics map[string]*memoryDatabaseTopic
|
|
}
|
|
|
|
type memoryDatabaseTopic struct {
|
|
messagesMutex sync.Mutex
|
|
messages []Message
|
|
}
|
|
|
|
func (t *memoryDatabaseTopic) addMessages(msgs []Message) error {
|
|
t.messagesMutex.Lock()
|
|
defer t.messagesMutex.Unlock()
|
|
if int64(len(t.messages)) != msgs[0].Offset {
|
|
return fmt.Errorf("message offset %d is not immediately after the previous offset %d", msgs[0].Offset, len(t.messages))
|
|
}
|
|
t.messages = append(t.messages, msgs...)
|
|
return nil
|
|
}
|
|
|
|
// getMessages returns the current messages as a slice.
|
|
// This slice will have it's own copy of the length field so won't be affected
|
|
// by adding more messages in addMessages.
|
|
// The slice will share the same backing array with the slice we append new
|
|
// messages to. It is safe to read the messages in the backing array since we
|
|
// only append to the slice. It is not safe to write or append to the returned
|
|
// slice.
|
|
func (t *memoryDatabaseTopic) getMessages() []Message {
|
|
t.messagesMutex.Lock()
|
|
defer t.messagesMutex.Unlock()
|
|
return t.messages
|
|
}
|
|
|
|
func (m *MemoryDatabase) getTopic(topicName string) *memoryDatabaseTopic {
|
|
m.topicsMutex.Lock()
|
|
defer m.topicsMutex.Unlock()
|
|
result := m.topics[topicName]
|
|
if result == nil {
|
|
result = &memoryDatabaseTopic{}
|
|
if m.topics == nil {
|
|
m.topics = map[string]*memoryDatabaseTopic{}
|
|
}
|
|
m.topics[topicName] = result
|
|
}
|
|
return result
|
|
}
|
|
|
|
// StoreMessages implements Database
|
|
func (m *MemoryDatabase) StoreMessages(topic string, messages []Message) error {
|
|
return m.getTopic(topic).addMessages(messages)
|
|
}
|
|
|
|
// FetchMessages implements Database
|
|
func (m *MemoryDatabase) FetchMessages(topic string, startOffset, endOffset int64) ([]Message, error) {
|
|
messages := m.getTopic(topic).getMessages()
|
|
if endOffset > int64(len(messages)) {
|
|
return nil, fmt.Errorf("end offset %d out of range %d", endOffset, len(messages))
|
|
}
|
|
if startOffset >= endOffset {
|
|
return nil, fmt.Errorf("start offset %d greater than or equal to end offset %d", startOffset, endOffset)
|
|
}
|
|
if startOffset < 0 {
|
|
return nil, fmt.Errorf("start offset %d less than 0", startOffset)
|
|
}
|
|
return messages[startOffset:endOffset], nil
|
|
}
|
|
|
|
// MaxOffsets implements Database
|
|
func (m *MemoryDatabase) MaxOffsets() (map[string]int64, error) {
|
|
m.topicsMutex.Lock()
|
|
defer m.topicsMutex.Unlock()
|
|
result := map[string]int64{}
|
|
for name, t := range m.topics {
|
|
result[name] = int64(len(t.getMessages())) - 1
|
|
}
|
|
return result, nil
|
|
}
|