├── .gitignore ├── LICENSE ├── README.md ├── README_CN.md ├── go.mod ├── helper.go ├── message.go ├── mq.go └── mq_test.go /.gitignore: -------------------------------------------------------------------------------- 1 | # If you prefer the allow list template instead of the deny list, see community template: 2 | # https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore 3 | # 4 | # Binaries for programs and plugins 5 | *.exe 6 | *.exe~ 7 | *.dll 8 | *.so 9 | *.dylib 10 | 11 | # Test binary, built with `go test -c` 12 | *.test 13 | 14 | # Output of the go coverage tool, specifically when used with LiteIDE 15 | *.out 16 | 17 | # Dependency directories (remove the comment below to include it) 18 | # vendor/ 19 | 20 | # Go workspace file 21 | go.work 22 | go.work.sum 23 | 24 | # env file 25 | .env 26 | *.dat 27 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 mobus 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | # mq 3 | 4 | [中文](README_CN.md) 5 | 6 | embed message queue for golang A high-performance, multi-producer, multi-consumer message queue library written in Go, based on memory-mapped files (`mmap`). This library supports maintaining message order, ensures successful message delivery, retains high performance, and avoids message loss. It also supports flexible features such as retaining recently consumed messages, periodic deletion of older consumed messages, and controlled message consumption rate. 7 | 8 | ## Features 9 | 10 | - **Memory-Mapped Files**: Efficient storage and retrieval of messages using memory-mapped files (`mmap`). 11 | - **Multi-Producer, Multi-Consumer**: Supports multiple producers and consumers simultaneously, ensuring thread safety and data integrity. 12 | - **Message Order**: Guarantees that messages are consumed in the order they are produced. 13 | - **No Message Loss**: Messages are persistently stored, ensuring no data loss even in case of service restarts. 14 | - **Control Message Consumption Rate**: Supports limiting message consumption speed (e.g., 30 messages per second). 15 | - **Service Restart Handling**: Automatically resumes from the last unconsumed message after service restarts. Also supports the option to re-consume messages from the beginning. 16 | - **New File Detection**: The library detects when the file is newly created and initializes offsets and metadata accordingly. 17 | 18 | Usage 19 | 20 | - Initialize the Message Queue 21 | 22 | ```go 23 | queue, err := NewMessageQueue("queue.dat", 1024*1024, 30) 24 | if err != nil { 25 | log.Fatalf("Failed to initialize the message queue: %v", err) 26 | } 27 | defer queue.Close() 28 | ``` 29 | 30 | Parameters: 31 | 32 | "queue.dat": The path to the message queue file. 33 | 1024*1024: The size of the message queue file (in bytes). 34 | 30: Message consumption rate limit (e.g., 30 messages per second). 35 | 10: Number of recently consumed messages to retain. 36 | 37 | - Produce Messages 38 | 39 | ```go 40 | msg := &Message{Data: []byte("Hello World")} 41 | err := queue.Push(msg) 42 | if err != nil { 43 | log.Fatalf("Failed to push message: %v", err) 44 | } 45 | ``` 46 | 47 | - Consume Messages 48 | 49 | ```go 50 | for { 51 | msg, err := queue.Pop() 52 | if err != nil { 53 | log.Fatalf("Failed to pop message: %v", err) 54 | } 55 | fmt.Printf("Consumed Message: %s\n", string(msg.Data)) 56 | } 57 | ``` 58 | 59 | Configuration 60 | Metadata Handling 61 | The library stores metadata, including writeOffset, readOffset, and isNewFile, at the head of the message queue file. This ensures that the library can resume operations from where it left off after service restarts. 62 | 63 | Consumption Rate Control 64 | The library uses a time.Ticker to control the message consumption rate. By setting the ratePerSecond parameter during initialization, you can control the number of messages consumed per second. Setting ratePerSecond to 0 disables rate limiting. 65 | 66 | Example 67 | 68 | ```go 69 | package main 70 | 71 | import ( 72 | "fmt" 73 | "log" 74 | "time" 75 | ) 76 | 77 | func main() { 78 | // Initialize the message queue with a rate limit of 30 messages per second and retain 10 consumed messages. 79 | queue, err := NewMessageQueue("queue.dat", 1024*1024, 30, 10) 80 | if err != nil { 81 | log.Fatalf("Failed to initialize the message queue: %v", err) 82 | } 83 | defer queue.Close() 84 | 85 | // Producer: Push messages into the queue 86 | go func() { 87 | for i := 0; i < 100; i++ { 88 | msg := &Message{Data: []byte(fmt.Sprintf("Message %d", i))} 89 | if err := queue.Push(msg); err != nil { 90 | log.Fatalf("Failed to push message: %v", err) 91 | } 92 | } 93 | }() 94 | 95 | // Consumer: Pop messages from the queue 96 | go func() { 97 | for { 98 | msg, err := queue.Pop() 99 | if err != nil { 100 | log.Fatalf("Failed to pop message: %v", err) 101 | } 102 | fmt.Printf("Consumed: %s\n", string(msg.Data)) 103 | } 104 | }() 105 | 106 | time.Sleep(5 * time.Second) 107 | queue.StopTicker() 108 | } 109 | ``` 110 | 111 | Contributing 112 | 113 | We welcome contributions! Here's how you can help: 114 | 115 | Fork the repository. 116 | Create a new branch for your feature or bug fix. 117 | Commit your changes and push them to your fork. 118 | Create a pull request with a detailed description of your changes. 119 | Testing 120 | To run the test cases, you can use the following command: 121 | 122 | ```bash 123 | go test ./... 124 | ``` 125 | 126 | We have included various test cases to validate the core functionalities such as: 127 | 128 | - Basic message push and pop operations. 129 | - Multi-producer and multi-consumer behavior. 130 | - Handling of rate-limited consumption. 131 | - Recovery from service restarts. 132 | 133 | License 134 | This project is licensed under the MIT License. See the LICENSE file for details. 135 | -------------------------------------------------------------------------------- /README_CN.md: -------------------------------------------------------------------------------- 1 | # 消息队列库 2 | 3 | 这是一个高性能的多生产者、多消费者消息队列库，基于 Go 语言开发，使用了内存映射文件（`mmap`）。该库支持保持消息顺序，确保消息成功投递，具有高性能，并且避免消息丢失。它还支持灵活的功能，例如保留最近消费的消息、定期删除已消费的旧消息，并支持控制消息消费速度。 4 | 5 | ## 功能 6 | 7 | - **内存映射文件**：使用内存映射文件（`mmap`）进行高效的消息存储和读取。 8 | - **多生产者、多消费者**：支持多个生产者和消费者同时操作，确保线程安全和数据一致性。 9 | - **消息顺序**：保证消息按照生产的顺序被消费。 10 | - **消息不丢失**：消息持久化存储，确保在服务重启后不会丢失数据。 11 | - **保留最近消费的消息**：可配置保留最后 N 条已消费的消息，并定期删除旧的已消费消息。 12 | - **控制消息消费速率**：支持限制消息消费的速度（例如每秒 30 条消息）。 13 | - **服务重启处理**：在服务重启后自动从未消费的消息继续消费，支持从头重新消费消息。 14 | - **新建文件检测**：库可以检测文件是否为新创建，必要时进行元数据初始化。 15 | 16 | 使用方法 17 | 18 | - 初始化消息队列 19 | 20 | ```go 21 | queue, err := NewMessageQueue("queue.dat", 1024*1024, 30, 10) 22 | if err != nil { 23 | log.Fatalf("无法初始化消息队列: %v", err) 24 | } 25 | defer queue.Close() 26 | ``` 27 | 28 | 参数： 29 | 30 | "queue.dat"：消息队列文件的路径。 31 | 1024*1024：消息队列文件的大小（以字节为单位）。 32 | 30：消息消费速率限制（例如每秒 30 条消息）。 33 | 10：保留最近消费的消息数量。 34 | 35 | - 生产消息 36 | 37 | ```go 38 | Copy code 39 | msg := &Message{Data: []byte("Hello World")} 40 | err := queue.Push(msg) 41 | if err != nil { 42 | log.Fatalf("推送消息失败: %v", err) 43 | } 44 | ``` 45 | 46 | - 消费消息 47 | 48 | ```go 49 | for { 50 | msg, err := queue.Pop() 51 | if err != nil { 52 | log.Fatalf("消费消息失败: %v", err) 53 | } 54 | fmt.Printf("消费消息: %s\n", string(msg.Data)) 55 | } 56 | ``` 57 | 58 | 配置说明 59 | 元数据处理 60 | 该库将元数据（包括 writeOffset、readOffset 和 isNewFile）存储在消息队列文件的头部。这确保了在服务重启后可以从上次未消费的位置继续运行。 61 | 62 | 消息消费速率控制 63 | 该库使用 time.Ticker 来控制消息消费速率。在初始化时，通过设置 ratePerSecond 参数可以控制每秒消费的消息数量。将 ratePerSecond 设置为 0 时，将禁用速率限制。 64 | 65 | 示例代码 66 | 67 | ```go 68 | package main 69 | 70 | import ( 71 | "fmt" 72 | "log" 73 | "time" 74 | ) 75 | 76 | func main() { 77 | // 初始化消息队列，消费速率限制为每秒 30 条消息，并保留最近 10 条消费的消息。 78 | queue, err := NewMessageQueue("queue.dat", 1024*1024, 30, 10) 79 | if err != nil { 80 | log.Fatalf("无法初始化消息队列: %v", err) 81 | } 82 | defer queue.Close() 83 | 84 | // 生产者：向队列中推送消息 85 | go func() { 86 | for i := 0; i < 100; i++ { 87 | msg := &Message{Data: []byte(fmt.Sprintf("Message %d", i))} 88 | if err := queue.Push(msg); err != nil { 89 | log.Fatalf("推送消息失败: %v", err) 90 | } 91 | } 92 | }() 93 | 94 | // 消费者：从队列中消费消息 95 | go func() { 96 | for { 97 | msg, err := queue.Pop() 98 | if err != nil { 99 | log.Fatalf("消费消息失败: %v", err) 100 | } 101 | fmt.Printf("消费消息: %s\n", string(msg.Data)) 102 | } 103 | }() 104 | 105 | time.Sleep(5 * 秒) 106 | queue.StopTicker() 107 | } 108 | ``` 109 | 110 | 贡献 111 | 欢迎贡献！以下是贡献代码的步骤： 112 | 113 | Fork 此仓库。 114 | 创建一个新分支用于你的功能或修复。 115 | 提交你的修改并推送到你的 Fork。 116 | 提交一个 Pull Request，详细描述你的更改。 117 | 测试 118 | 使用以下命令运行测试： 119 | 120 | ```bash 121 | go test ./... 122 | ``` 123 | 124 | 我们的测试用例涵盖了以下核心功能： 125 | 126 | - 基本的消息推送和消费操作。 127 | - 多生产者和多消费者的行为。 128 | - 消息消费速率限制的处理。 129 | - 服务重启后的恢复能力。 130 | 131 | 许可证 132 | 此项目基于 MIT 许可证。请查看 LICENSE 文件以获取详细信息。 133 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/sunvim/mq 2 | 3 | go 1.23.2 4 | -------------------------------------------------------------------------------- /helper.go: -------------------------------------------------------------------------------- 1 | package mq 2 | 3 | import ( 4 | "encoding/binary" 5 | "os" 6 | ) 7 | 8 | // saveMetadata writes metadata to the file header 9 | func (mq *MessageQueue) saveMetadata() error { 10 | mq.metadataMux.Lock() 11 | defer mq.metadataMux.Unlock() 12 | 13 | buf := make([]byte, MetadataSize) 14 | binary.LittleEndian.PutUint64(buf[0:], uint64(mq.metadata.WriteOffset)) 15 | binary.LittleEndian.PutUint64(buf[8:], uint64(mq.metadata.ReadOffset)) 16 | 17 | _, err := mq.file.WriteAt(buf, 0) // write to the file header 18 | if err != nil { 19 | return err 20 | } 21 | 22 | return mq.file.Sync() // ensure data is written to disk 23 | } 24 | 25 | // loadMetadata loads metadata from the file header 26 | func (mq *MessageQueue) loadMetadata() error { 27 | buf := make([]byte, MetadataSize) 28 | _, err := mq.file.ReadAt(buf, 0) 29 | if err != nil { 30 | return err 31 | } 32 | 33 | mq.metadataMux.Lock() 34 | defer mq.metadataMux.Unlock() 35 | 36 | mq.metadata.WriteOffset = int64(binary.LittleEndian.Uint64(buf[0:])) 37 | mq.metadata.ReadOffset = int64(binary.LittleEndian.Uint64(buf[8:])) 38 | 39 | return nil 40 | } 41 | 42 | // showMetadata displays the metadata in the file header 43 | func (mq *MessageQueue) showMetadata() QueueMetadata { 44 | mq.metadataMux.RLock() 45 | defer mq.metadataMux.RUnlock() 46 | 47 | return mq.metadata 48 | } 49 | 50 | // consumeMessage is a helper function to actually consume a message 51 | func (mq *MessageQueue) consumeMessage() (*Message, error) { 52 | mq.readMutex.Lock() 53 | defer mq.readMutex.Unlock() 54 | 55 | // Get current read and write offsets 56 | readOffset := mq.getReadOffset() 57 | writeOffset := mq.getWriteOffset() 58 | 59 | // If there are no messages to read, wait for new ones 60 | if readOffset >= writeOffset { 61 | mq.readMutex.Unlock() // Release lock before waiting 62 | mq.WaitForMessages() // This will block until new messages are available 63 | mq.readMutex.Lock() // Re-acquire lock 64 | 65 | // Re-read offsets after wait 66 | readOffset = mq.getReadOffset() 67 | writeOffset = mq.getWriteOffset() 68 | 69 | // Double-check there are messages after waiting 70 | if readOffset >= writeOffset { 71 | return nil, os.ErrNotExist 72 | } 73 | } 74 | 75 | // read message length 76 | msgLen := int(binary.LittleEndian.Uint32(mq.mmap[readOffset:])) 77 | readOffset += 4 78 | 79 | if readOffset+int64(msgLen) > mq.size { 80 | return nil, os.ErrInvalid // out of range 81 | } 82 | 83 | // read message content 84 | msgData := make([]byte, msgLen) 85 | copy(msgData, mq.mmap[readOffset:readOffset+int64(msgLen)]) 86 | readOffset += int64(msgLen) 87 | 88 | // Update read offset 89 | mq.setReadOffset(readOffset) 90 | 91 | return &Message{Data: msgData}, nil 92 | } 93 | -------------------------------------------------------------------------------- /message.go: -------------------------------------------------------------------------------- 1 | package mq 2 | 3 | import ( 4 | "bytes" 5 | "encoding/binary" 6 | ) 7 | 8 | // Message is a variable-length message structure with an added length field 9 | type Message struct { 10 | Length int32 // Message length 11 | Data []byte // Actual message data 12 | } 13 | 14 | // Encode encodes the message into a byte array 15 | func (msg *Message) Encode() ([]byte, error) { 16 | buf := new(bytes.Buffer) 17 | 18 | msg.Length = int32(len(msg.Data)) 19 | 20 | // First write the message length 21 | if err := binary.Write(buf, binary.LittleEndian, msg.Length); err != nil { 22 | return nil, err 23 | } 24 | 25 | // Write the message data 26 | if _, err := buf.Write(msg.Data); err != nil { 27 | return nil, err 28 | } 29 | 30 | return buf.Bytes(), nil 31 | } 32 | 33 | // Decode decodes the message from a byte array 34 | func Decode(data []byte) (*Message, error) { 35 | buf := bytes.NewReader(data) 36 | 37 | // Read the message length 38 | var length int32 39 | if err := binary.Read(buf, binary.LittleEndian, &length); err != nil { 40 | return nil, err 41 | } 42 | 43 | // Read the message data 44 | msgData := make([]byte, length) 45 | if _, err := buf.Read(msgData); err != nil { 46 | return nil, err 47 | } 48 | 49 | return &Message{ 50 | Length: length, 51 | Data: msgData, 52 | }, nil 53 | } 54 | -------------------------------------------------------------------------------- /mq.go: -------------------------------------------------------------------------------- 1 | package mq 2 | 3 | import ( 4 | "encoding/binary" 5 | "os" 6 | "sync" 7 | "syscall" 8 | "time" 9 | ) 10 | 11 | const ( 12 | MetadataSize = 128 // Reserve 128 bytes at the beginning of the file for storing metadata 13 | ) 14 | 15 | // QueueMetadata defines the metadata structure for the message queue 16 | type QueueMetadata struct { 17 | WriteOffset int64 // Write offset 18 | ReadOffset int64 // Read offset 19 | } 20 | 21 | // MessageQueue is the main structure of the message queue 22 | type MessageQueue struct { 23 | file *os.File // File mapped to memory 24 | mmap []byte // Mapped memory 25 | size int64 // Size of the mapped file 26 | metadata QueueMetadata // Message queue metadata 27 | writeMutex sync.Mutex // Protect multi-producer writes 28 | readMutex sync.Mutex // Protect multi-consumer reads 29 | cond *sync.Cond // Synchronization between producers and consumers 30 | ticker *time.Ticker // Control consumption frequency 31 | metadataMux sync.RWMutex // Protects metadata access 32 | } 33 | 34 | // NewMessageQueue creates a new message queue and maps its file to memory 35 | func NewMessageQueue(filePath string, size int64, ratePerSecond int) (*MessageQueue, error) { 36 | file, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE, 0644) 37 | if err != nil { 38 | return nil, err 39 | } 40 | 41 | // Pre-allocate file size 42 | if err := file.Truncate(size); err != nil { 43 | return nil, err 44 | } 45 | 46 | // Use mmap to map the file to memory 47 | mmap, err := syscall.Mmap(int(file.Fd()), 0, int(size), syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_SHARED) 48 | if err != nil { 49 | return nil, err 50 | } 51 | 52 | // Initialize ticker to control consumption frequency 53 | var ticker *time.Ticker 54 | if ratePerSecond > 0 { 55 | ticker = time.NewTicker(time.Second / time.Duration(ratePerSecond)) 56 | } 57 | 58 | mq := &MessageQueue{ 59 | file: file, 60 | mmap: mmap, 61 | size: size, 62 | ticker: ticker, 63 | } 64 | mq.cond = sync.NewCond(&mq.writeMutex) 65 | 66 | // Load metadata 67 | if err := mq.loadMetadata(); err != nil { 68 | return nil, err 69 | } 70 | 71 | if mq.metadata.WriteOffset < MetadataSize { 72 | mq.metadata.WriteOffset = MetadataSize 73 | } 74 | if mq.metadata.ReadOffset < MetadataSize { 75 | mq.metadata.ReadOffset = MetadataSize 76 | } 77 | 78 | return mq, nil 79 | } 80 | 81 | // Close closes the message queue and cleans up resources 82 | func (mq *MessageQueue) Close() error { 83 | err := mq.saveMetadata() 84 | if err != nil { 85 | return err 86 | } 87 | 88 | if err := syscall.Munmap(mq.mmap); err != nil { 89 | return err 90 | } 91 | return mq.file.Close() 92 | } 93 | 94 | // Push writes a message to the queue 95 | func (mq *MessageQueue) Push(msg *Message) error { 96 | mq.writeMutex.Lock() 97 | defer mq.writeMutex.Unlock() 98 | 99 | // Get current write offset 100 | currentWriteOffset := mq.getWriteOffset() 101 | 102 | dataLen := len(msg.Data) 103 | if currentWriteOffset+int64(dataLen+4) > mq.size { 104 | return os.ErrInvalid // Queue is full 105 | } 106 | 107 | // Write message length 108 | binary.LittleEndian.PutUint32(mq.mmap[currentWriteOffset:], uint32(dataLen)) 109 | currentWriteOffset += 4 110 | 111 | // Write message data 112 | copy(mq.mmap[currentWriteOffset:], msg.Data) 113 | currentWriteOffset += int64(dataLen) 114 | 115 | // Update write offset 116 | mq.setWriteOffset(currentWriteOffset) 117 | 118 | // Notify consumers of new message 119 | mq.cond.Signal() 120 | 121 | return nil 122 | } 123 | 124 | // Pop reads a variable-length message from the queue and consumes it at the set frequency 125 | func (mq *MessageQueue) Pop() (*Message, error) { 126 | for { 127 | if mq.ticker == nil { // Consume immediately if there is no frequency limit 128 | return mq.consumeMessage() 129 | } 130 | 131 | select { 132 | case <-mq.ticker.C: // Consume messages at the set frequency 133 | return mq.consumeMessage() 134 | } 135 | } 136 | } 137 | 138 | // PopAll reads all unconsumed messages from the queue at once 139 | // This returns all available messages without waiting for new ones 140 | func (mq *MessageQueue) PopAll() ([]*Message, error) { 141 | mq.readMutex.Lock() 142 | defer mq.readMutex.Unlock() 143 | 144 | // Get current read and write offsets 145 | readOffset := mq.getReadOffset() 146 | writeOffset := mq.getWriteOffset() 147 | 148 | // If there are no messages to read 149 | if readOffset >= writeOffset { 150 | return []*Message{}, nil 151 | } 152 | 153 | var messages []*Message 154 | currentOffset := readOffset 155 | 156 | // Read all available messages 157 | for currentOffset < writeOffset { 158 | // Check if we have enough space to read message length 159 | if currentOffset+4 > writeOffset { 160 | break 161 | } 162 | 163 | // Read message length 164 | msgLen := int(binary.LittleEndian.Uint32(mq.mmap[currentOffset:])) 165 | currentOffset += 4 166 | 167 | // Check if the complete message is available 168 | if currentOffset+int64(msgLen) > writeOffset { 169 | break 170 | } 171 | 172 | // Read message content 173 | msgData := make([]byte, msgLen) 174 | copy(msgData, mq.mmap[currentOffset:currentOffset+int64(msgLen)]) 175 | currentOffset += int64(msgLen) 176 | 177 | messages = append(messages, &Message{Data: msgData}) 178 | } 179 | 180 | // Update read offset 181 | mq.setReadOffset(currentOffset) 182 | 183 | return messages, nil 184 | } 185 | 186 | // WaitForMessages waits until there are new messages available 187 | func (mq *MessageQueue) WaitForMessages() { 188 | mq.writeMutex.Lock() 189 | for mq.getReadOffset() >= mq.getWriteOffset() { 190 | mq.cond.Wait() // wait for new messages 191 | } 192 | mq.writeMutex.Unlock() 193 | } 194 | 195 | // Flush flushes the queue, persisting messages in memory to disk 196 | func (mq *MessageQueue) Flush() error { 197 | mq.writeMutex.Lock() 198 | defer mq.writeMutex.Unlock() 199 | 200 | return mq.file.Sync() 201 | } 202 | 203 | // DeleteConsumedMessages deletes consumed messages and updates the queue offsets 204 | func (mq *MessageQueue) DeleteConsumedMessages() error { 205 | // Lock both read and write to perform this operation 206 | mq.writeMutex.Lock() 207 | mq.readMutex.Lock() 208 | defer mq.readMutex.Unlock() 209 | defer mq.writeMutex.Unlock() 210 | 211 | readOffset := mq.getReadOffset() 212 | writeOffset := mq.getWriteOffset() 213 | 214 | // If there are no consumed messages, return directly 215 | if readOffset == MetadataSize { 216 | return nil // No consumed messages 217 | } 218 | 219 | // Calculate the size of unconsumed messages 220 | remainingDataSize := writeOffset - readOffset 221 | if remainingDataSize > 0 { 222 | // Move unconsumed messages to the beginning of the queue 223 | copy(mq.mmap[MetadataSize:], mq.mmap[readOffset:writeOffset]) 224 | } 225 | 226 | // Update offsets 227 | mq.setWriteOffset(MetadataSize + remainingDataSize) 228 | mq.setReadOffset(MetadataSize) 229 | 230 | return nil 231 | } 232 | 233 | // Helper methods for thread-safe metadata access 234 | 235 | func (mq *MessageQueue) getReadOffset() int64 { 236 | mq.metadataMux.RLock() 237 | defer mq.metadataMux.RUnlock() 238 | return mq.metadata.ReadOffset 239 | } 240 | 241 | func (mq *MessageQueue) getWriteOffset() int64 { 242 | mq.metadataMux.RLock() 243 | defer mq.metadataMux.RUnlock() 244 | return mq.metadata.WriteOffset 245 | } 246 | 247 | func (mq *MessageQueue) setReadOffset(offset int64) { 248 | mq.metadataMux.Lock() 249 | defer mq.metadataMux.Unlock() 250 | mq.metadata.ReadOffset = offset 251 | } 252 | 253 | func (mq *MessageQueue) setWriteOffset(offset int64) { 254 | mq.metadataMux.Lock() 255 | defer mq.metadataMux.Unlock() 256 | mq.metadata.WriteOffset = offset 257 | } 258 | -------------------------------------------------------------------------------- /mq_test.go: -------------------------------------------------------------------------------- 1 | package mq 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | "testing" 7 | "time" 8 | ) 9 | 10 | func TestPopAll(t *testing.T) { 11 | // Initialize the message queue 12 | queue, err := NewMessageQueue("test_queue_popall.dat", 1024*1024, 0) 13 | if err != nil { 14 | t.Fatalf("failed to initialize message queue: %v", err) 15 | } 16 | defer queue.Close() 17 | 18 | // Test with empty queue 19 | msgs, err := queue.PopAll() 20 | if err != nil { 21 | t.Fatalf("failed to pop all messages: %v", err) 22 | } 23 | if len(msgs) != 0 { 24 | t.Fatalf("expected 0 messages, got %d", len(msgs)) 25 | } 26 | 27 | // Produce 50 messages 28 | for i := range 50 { 29 | buf := []byte("Message") 30 | msg := &Message{Data: fmt.Appendf(buf, "%d", i)} 31 | if err := queue.Push(msg); err != nil { 32 | t.Fatalf("failed to push message: %v", err) 33 | } 34 | } 35 | 36 | // Pop all messages at once 37 | msgs, err = queue.PopAll() 38 | if err != nil { 39 | t.Fatalf("failed to pop all messages: %v", err) 40 | } 41 | if len(msgs) != 50 { 42 | t.Fatalf("expected 50 messages, got %d", len(msgs)) 43 | } 44 | 45 | // Verify message contents 46 | for i, msg := range msgs { 47 | if string(msg.Data) != fmt.Sprintf("Message%d", i) { 48 | t.Fatalf("message order incorrect, expected: Message %d, got: %s", i, msg.Data) 49 | } 50 | } 51 | 52 | // Verify that all messages were consumed 53 | msgs, err = queue.PopAll() 54 | if err != nil { 55 | t.Fatalf("failed to pop all messages: %v", err) 56 | } 57 | if len(msgs) != 0 { 58 | t.Fatalf("expected 0 messages after consuming all, got %d", len(msgs)) 59 | } 60 | } 61 | 62 | func TestConcurrentReadWrite(t *testing.T) { 63 | // Initialize the message queue 64 | queue, err := NewMessageQueue("test_queue_concurrent.dat", 1024*1024, 0) 65 | if err != nil { 66 | t.Fatalf("failed to initialize message queue: %v", err) 67 | } 68 | defer queue.Close() 69 | 70 | // Number of messages to produce/consume 71 | const messageCount = 1000 72 | // Channel to signal completion 73 | done := make(chan bool) 74 | 75 | // Producer goroutine 76 | go func() { 77 | for i := range messageCount { 78 | buf := []byte("Message") 79 | msg := &Message{Data: fmt.Appendf(buf, "%d", i)} 80 | if err := queue.Push(msg); err != nil { 81 | t.Errorf("producer: failed to push message: %v", err) 82 | return 83 | } 84 | // Small sleep to make concurrent issues more likely to show up 85 | time.Sleep(time.Millisecond) 86 | } 87 | }() 88 | 89 | // Consumer goroutine - using PopAll 90 | go func() { 91 | consumed := 0 92 | for consumed < messageCount { 93 | msgs, err := queue.PopAll() 94 | if err != nil { 95 | t.Errorf("consumer: failed to pop messages: %v", err) 96 | return 97 | } 98 | consumed += len(msgs) 99 | // Small sleep to make concurrent issues more likely to show up 100 | time.Sleep(2 * time.Millisecond) 101 | } 102 | done <- true 103 | }() 104 | 105 | // Wait for consumer to finish 106 | select { 107 | case <-done: 108 | // Success case 109 | case <-time.After(10 * time.Second): 110 | t.Fatal("test timed out") 111 | } 112 | } 113 | 114 | func TestBasicPushAndPop(t *testing.T) { 115 | queue, err := NewMessageQueue("test_queue_basic.dat", 1024*1024, 0) 116 | if err != nil { 117 | t.Fatalf("failed to create message queue: %v", err) 118 | } 119 | defer queue.Close() 120 | 121 | msg1 := &Message{Data: []byte("Hello")} 122 | msg2 := &Message{Data: []byte("World")} 123 | 124 | // Write messages 125 | err = queue.Push(msg1) 126 | if err != nil { 127 | t.Fatalf("failed to push message: %v", err) 128 | } 129 | 130 | err = queue.Push(msg2) 131 | if err != nil { 132 | t.Fatalf("failed to push message: %v", err) 133 | } 134 | 135 | // Read messages 136 | poppedMsg1, err := queue.Pop() 137 | if err != nil { 138 | t.Fatalf("failed to pop message: %v", err) 139 | } 140 | if string(poppedMsg1.Data) != "Hello" { 141 | t.Fatalf("expected 'Hello', got '%s'", string(poppedMsg1.Data)) 142 | } 143 | 144 | poppedMsg2, err := queue.Pop() 145 | if err != nil { 146 | t.Fatalf("failed to pop message: %v", err) 147 | } 148 | if string(poppedMsg2.Data) != "World" { 149 | t.Fatalf("expected 'World', got '%s'", string(poppedMsg2.Data)) 150 | } 151 | } 152 | 153 | func TestMessageOrder(t *testing.T) { 154 | queue, err := NewMessageQueue("test_queue_order.dat", 1024*1024, 0) 155 | if err != nil { 156 | t.Fatalf("failed to create message queue: %v", err) 157 | } 158 | defer queue.Close() 159 | 160 | messages := []string{"Msg1", "Msg2", "Msg3"} 161 | 162 | // Write multiple messages 163 | for _, msg := range messages { 164 | err := queue.Push(&Message{Data: []byte(msg)}) 165 | if err != nil { 166 | t.Fatalf("failed to push message: %v", err) 167 | } 168 | } 169 | 170 | // Read and verify order 171 | for _, expected := range messages { 172 | msg, err := queue.Pop() 173 | if err != nil { 174 | t.Fatalf("failed to pop message: %v", err) 175 | } 176 | if string(msg.Data) != expected { 177 | t.Fatalf("expected '%s', got '%s'", expected, string(msg.Data)) 178 | } 179 | } 180 | } 181 | 182 | func TestFrequencyLimit(t *testing.T) { 183 | queue, err := NewMessageQueue("test_queue_freq.dat", 1024*1024, 2) // Consume at most 2 messages per second 184 | if err != nil { 185 | t.Fatalf("failed to create message queue: %v", err) 186 | } 187 | defer queue.Close() 188 | 189 | for range 5 { 190 | err := queue.Push(&Message{Data: []byte("Message")}) 191 | if err != nil { 192 | t.Fatalf("failed to push message: %v", err) 193 | } 194 | } 195 | 196 | start := time.Now() 197 | 198 | // Consume 5 messages 199 | for range 5 { 200 | _, err := queue.Pop() 201 | if err != nil { 202 | t.Fatalf("failed to pop message: %v", err) 203 | } 204 | } 205 | 206 | elapsed := time.Since(start) 207 | 208 | // Expect to consume 5 messages in about 2.5 seconds (since the rate is 2 messages per second) 209 | if elapsed < 2*time.Second || elapsed > 3*time.Second { 210 | t.Fatalf("expected consumption time around 2.5 seconds, got %v", elapsed) 211 | } 212 | } 213 | 214 | func TestNoFrequencyLimit(t *testing.T) { 215 | queue, err := NewMessageQueue("test_queue_no_freq.dat", 1024*1024, 0) // No frequency limit 216 | if err != nil { 217 | t.Fatalf("failed to create message queue: %v", err) 218 | } 219 | defer queue.Close() 220 | 221 | for range 10 { 222 | err := queue.Push(&Message{Data: []byte("Message")}) 223 | if err != nil { 224 | t.Fatalf("failed to push message: %v", err) 225 | } 226 | } 227 | 228 | start := time.Now() 229 | 230 | // Consume 10 messages 231 | for range 10 { 232 | _, err := queue.Pop() 233 | if err != nil { 234 | t.Fatalf("failed to pop message: %v", err) 235 | } 236 | } 237 | 238 | elapsed := time.Since(start) 239 | 240 | // Since there is no frequency limit, the expected time is very short 241 | if elapsed > 100*time.Millisecond { 242 | t.Fatalf("expected fast consumption, but got %v", elapsed) 243 | } 244 | } 245 | 246 | func TestMultiProducerConsumer(t *testing.T) { 247 | queue, err := NewMessageQueue("test_queue_multi.dat", 1024*1024, 0) 248 | if err != nil { 249 | t.Fatalf("failed to create message queue: %v", err) 250 | } 251 | defer queue.Close() 252 | 253 | totalMessages := 100 254 | produced := 0 255 | consumed := 0 256 | var mu sync.Mutex 257 | errCh := make(chan error, 10) // Channel for error transmission 258 | 259 | // Multiple producers 260 | producer := func() { 261 | for range totalMessages / 2 { 262 | err := queue.Push(&Message{Data: []byte("Message")}) 263 | if err != nil { 264 | errCh <- err 265 | return 266 | } 267 | mu.Lock() 268 | produced++ 269 | mu.Unlock() 270 | } 271 | } 272 | 273 | // Multiple consumers 274 | consumer := func() { 275 | for { 276 | _, err := queue.Pop() 277 | if err != nil { 278 | errCh <- err 279 | return 280 | } 281 | mu.Lock() 282 | consumed++ 283 | mu.Unlock() 284 | } 285 | } 286 | 287 | // Start 2 producers and 2 consumers 288 | go producer() 289 | go producer() 290 | go consumer() 291 | go consumer() 292 | 293 | // Wait for a while to simulate the production and consumption process 294 | time.Sleep(2 * time.Second) 295 | 296 | // Check for errors 297 | select { 298 | case err := <-errCh: 299 | if err != nil { 300 | t.Fatalf("error occurred during test: %v", err) 301 | } 302 | default: 303 | // If no errors, continue to check the production and consumption counts 304 | } 305 | 306 | if produced != totalMessages { 307 | t.Fatalf("expected %d messages produced, but got %d", totalMessages, produced) 308 | } 309 | 310 | if consumed != totalMessages { 311 | t.Fatalf("expected %d messages consumed, but got %d", totalMessages, consumed) 312 | } 313 | } 314 | 315 | func TestDeleteConsumedMessages(t *testing.T) { 316 | // Initialize the message queue 317 | queue, err := NewMessageQueue("test_queue_delete_consumed.dat", 1024*1024, 0) 318 | if err != nil { 319 | t.Fatalf("failed to initialize message queue: %v", err) 320 | } 321 | defer queue.Close() 322 | 323 | // Produce 50 messages 324 | for i := range 50 { 325 | buf := []byte("Message") 326 | msg := &Message{Data: fmt.Appendf(buf, "%d", i)} 327 | if err := queue.Push(msg); err != nil { 328 | t.Fatalf("failed to push message: %v", err) 329 | } 330 | } 331 | 332 | // Consume the first 30 messages 333 | for i := range 30 { 334 | msg, err := queue.Pop() 335 | if err != nil { 336 | t.Fatalf("failed to pop message: %v", err) 337 | } 338 | if string(msg.Data) != fmt.Sprintf("Message%d", i) { 339 | t.Fatalf("message order incorrect, expected: %d, got: %s", i, msg.Data) 340 | } 341 | } 342 | 343 | // Call DeleteConsumedMessages to delete the consumed 30 messages 344 | err = queue.DeleteConsumedMessages() 345 | if err != nil { 346 | t.Fatalf("failed to delete consumed messages: %v", err) 347 | } 348 | 349 | // Verify that the remaining messages are correctly retained 350 | for i := 30; i < 50; i++ { 351 | msg, err := queue.Pop() 352 | if err != nil { 353 | t.Fatalf("failed to pop message: %v", err) 354 | } 355 | if string(msg.Data) != fmt.Sprintf("Message%d", i) { 356 | t.Fatalf("message order incorrect, expected: %d, got: %s", i, msg.Data) 357 | } 358 | } 359 | 360 | } 361 | --------------------------------------------------------------------------------