特性：支持 commit 读取进度及 close 优化

Author: yoonper

README.md

@@ -14,34 +14,55 @@ Disk-based FIFO queue
 ## Getting Started
 
 ```
-go get -u github.com/4x99/diskqueue
+go get -u github.com/yoonper/diskqueue
 ```
 
 ```
 package main
 
 import (
 	"fmt"
-	"github.com/4x99/diskqueue"
+	"github.com/yoonper/diskqueue"
 	"log"
+	"time"
 )
 
 func main() {
+	var err error
+	var queue *diskqueue.Diskqueue
+
+	// config
+	diskqueue.Config.Path = "/tmp/diskqueue"
+	diskqueue.Config.BatchSize = 1
+
 	// start
-	diskqueue.Config.Path = "/tmp"
-	queue, err := diskqueue.Start()
-	if err != nil {
+	if queue, err = diskqueue.Start(); err != nil {
 		log.Fatalln(err)
 	}
 
-	// write data
-	err = queue.Write([]byte("data"))
-	fmt.Println(err)
+	// write
+	go func() {
+		for {
+			time.Sleep(time.Second)
+			data := []byte(time.Now().Format("2006-01-02 15:04:05"))
+			if err := queue.Write(data); err != nil {
+				fmt.Println(err)
+			}
+		}
+	}()
 
-	// read data
-	if data, err := queue.Read(); err != nil {
-		fmt.Println(data)
-	}
+	// read
+	go func() {
+		for {
+			time.Sleep(time.Second)
+			if index, offset, data, err := queue.Read(); err == nil {
+				fmt.Println(index, offset, string(data))
+				queue.Commit(index, offset) // commit
+			}
+		}
+	}()
+
+	select {}
 }
 ```
 

config.go

@@ -12,6 +12,6 @@ type config struct {
 	BatchTime        time.Duration // interval per sync
 	SegmentSize      int64         // size of each segment (in bytes)
 	SegmentLimit     int64         // max number of segment
-	CheckpointFile   string        // record read offset
+	CheckpointFile   string        // read index and offset
 	MinRequiredSpace int64         // minimum required free space (in bytes)
 }

diskqueue.go

@@ -1,6 +1,7 @@
 package diskqueue
 
 import (
+	"context"
 	"errors"
 	"io"
 	"os"
@@ -10,9 +11,11 @@ import (
 
 type Diskqueue struct {
 	sync.RWMutex
-	close     bool
-	closeChan chan bool
-	ticker    *time.Ticker
+	close  bool
+	ticker *time.Ticker
+	wg     *sync.WaitGroup
+	ctx    context.Context
+	cancel context.CancelFunc
 }
 
 var (
@@ -36,10 +39,10 @@ func Start() (*Diskqueue, error) {
 	if _, err := os.Stat(Config.Path); err != nil {
 		return nil, err
 	}
-
-	queue := &Diskqueue{close: false, closeChan: make(chan bool)}
+	queue := &Diskqueue{close: false, wg: &sync.WaitGroup{}}
 	queue.ticker = time.NewTicker(Config.BatchTime)
-	Reader.restore()
+	queue.ctx, queue.cancel = context.WithCancel(context.TODO())
+	_ = Reader.restore()
 	go queue.sync()
 	return queue, nil
 }
@@ -57,19 +60,30 @@ func (queue *Diskqueue) Write(data []byte) error {
 }
 
 // Read data
-func (queue *Diskqueue) Read() ([]byte, error) {
+func (queue *Diskqueue) Read() (int64, int64, []byte, error) {
 	if queue.close {
-		return nil, errors.New("closed")
+		return 0, 0, nil, errors.New("closed")
 	}
 
 	queue.RLock()
 	defer queue.RUnlock()
 
-	data, err := Reader.read()
+	index, offset, data, err := Reader.read()
 	if err == io.EOF && (Writer.file == nil || Reader.file.Name() != Writer.file.Name()) {
 		_ = Reader.rotate()
 	}
-	return data, err
+	return index, offset, data, err
+}
+
+// Commit index and offset
+func (queue *Diskqueue) Commit(index int64, offset int64) {
+	if queue.close {
+		return
+	}
+
+	ck := &Reader.checkpoint
+	ck.Index, ck.Offset = index, offset
+	Reader.sync()
 }
 
 // Close diskqueue
@@ -79,21 +93,23 @@ func (queue *Diskqueue) Close() {
 	}
 
 	queue.close = true
-	queue.closeChan <- true
+	queue.cancel()
+	queue.wg.Wait()
 	Writer.close()
 	Reader.close()
 }
 
 // sync data
 func (queue *Diskqueue) sync() {
+	queue.wg.Add(1)
+	defer queue.wg.Done()
 	for {
 		select {
 		case <-queue.ticker.C:
 			queue.Lock()
 			Writer.sync()
-			Reader.sync()
 			queue.Unlock()
-		case <-queue.closeChan:
+		case <-queue.ctx.Done():
 			return
 		}
 	}

go.mod

@@ -1,3 +1,3 @@
 module github.com/4x99/diskqueue
 
-go 1.17
+go 1.18

reader.go

@@ -2,7 +2,9 @@ package diskqueue
 
 import (
 	"bufio"
+	"encoding/json"
 	"errors"
+	"fmt"
 	"io/ioutil"
 	"os"
 	"path"
@@ -12,51 +14,63 @@ import (
 )
 
 type reader struct {
-	file   *os.File
-	offset int64
-	reader *bufio.Reader
+	file       *os.File
+	index      int64
+	offset     int64
+	reader     *bufio.Reader
+	checkpoint checkpoint
+}
+
+type checkpoint struct {
+	Index  int64 `json:"index"`
+	Offset int64 `json:"offset"`
 }
 
 // read data
-func (r *reader) read() ([]byte, error) {
-	// open a new segment
-	if err := r.open(); err != nil {
-		return nil, err
+func (r *reader) read() (int64, int64, []byte, error) {
+	if err := r.check(); err != nil {
+		return r.index, r.offset, nil, err
 	}
 
 	// read a line
 	data, _, err := r.reader.ReadLine()
 	if err != nil {
-		return nil, err
+		return r.index, r.offset, nil, err
 	}
 
 	r.offset += int64(len(data)) + 1
-	return data, err
+	return r.index, r.offset, data, err
 }
 
-// open a new segment
-func (r *reader) open() error {
+// check a new segment
+func (r *reader) check() error {
 	if r.file != nil {
 		return nil
 	}
 
-	files, err := r.list()
+	file, err := r.next()
 	if err != nil {
 		return err
 	}
 
-	// open the earliest segment
-	if r.file, err = os.OpenFile(files[0], os.O_RDONLY, Config.FilePerm); err != nil {
+	return r.open(file)
+}
+
+func (r *reader) open(file string) (err error) {
+	if r.file, err = os.OpenFile(file, os.O_RDONLY, Config.FilePerm); err != nil {
 		return err
 	}
 
+	// get file index
+	r.index = r.getIndex(file)
+
 	// seek read offset
 	if _, err = r.file.Seek(r.offset, 0); err != nil {
 		return err
 	}
 
 	r.reader = bufio.NewReader(r.file)
-	return err
+	return nil
 }
 
 // rotate to next segment
@@ -66,58 +80,76 @@ func (r *reader) rotate() error {
 	}
 
 	// close segment
-	r.file.Close()
-
-	// remove segment
-	if err := os.Remove(r.file.Name()); err != nil {
-		return err
-	}
-
+	_ = r.file.Close()
 	r.file, r.offset, r.reader = nil, 0, nil
-	r.sync()
-
 	return nil
 }
 
-// list all segments
-func (r *reader) list() ([]string, error) {
-	files, err := filepath.Glob(filepath.Join(Config.Path, "*.data"))
-	if err != nil {
-		return nil, err
+// close reader
+func (r *reader) close() {
+	if r.file == nil {
+		return
 	}
 
-	if len(files) == 0 {
-		return nil, errors.New("empty")
+	if err := r.file.Close(); err != nil {
+		return
 	}
 
-	sort.Strings(files)
-	return files, err
+	r.file, r.reader, r.index, r.offset = nil, nil, 0, 0
 }
 
-// sync read offset
+// sync index and offset
 func (r *reader) sync() {
 	name := path.Join(Config.Path, Config.CheckpointFile)
-	offset := []byte(strconv.FormatInt(r.offset, 10))
-	_ = ioutil.WriteFile(name, offset, Config.FilePerm)
+	data, _ := json.Marshal(&r.checkpoint)
+	_ = ioutil.WriteFile(name, data, Config.FilePerm)
 }
 
-// close reader
-func (r *reader) close() {
-	if r.file == nil {
-		return
+// restore index and offset
+func (r *reader) restore() (err error) {
+	name := path.Join(Config.Path, Config.CheckpointFile)
+	data, _ := ioutil.ReadFile(name)
+	_ = json.Unmarshal(data, &r.checkpoint)
+	r.index, r.offset = r.checkpoint.Index, r.checkpoint.Offset
+
+	if r.index == 0 {
+		return nil
 	}
 
-	r.sync()
-	if err := r.file.Close(); err != nil {
-		return
+	if err = r.open(fmt.Sprintf("%s/%d.data", Config.Path, r.index)); err != nil {
+		r.offset = 0
 	}
 
-	r.file, r.offset, r.reader = nil, 0, nil
+	return err
 }
 
-// restore read offset
-func (r *reader) restore() {
-	name := path.Join(Config.Path, Config.CheckpointFile)
-	offset, _ := ioutil.ReadFile(name)
-	r.offset, _ = strconv.ParseInt(string(offset), 10, 64)
+// next segment
+func (r *reader) next() (string, error) {
+	files, err := filepath.Glob(filepath.Join(Config.Path, "*.data"))
+	if err != nil {
+		return "", err
+	}
+
+	sort.Strings(files)
+
+	for _, file := range files {
+		index := r.getIndex(file)
+		if index < r.checkpoint.Index {
+			_ = os.Remove(file) // remove expired segment
+		}
+
+		if index > r.index {
+			return file, nil
+		}
+	}
+
+	return "", errors.New("queue is empty")
+}
+
+// get segment index
+func (r *reader) getIndex(filename string) int64 {
+	base := path.Base(filename)
+	name := base[0 : len(base)-len(path.Ext(filename))]
+	index, _ := strconv.ParseInt(name, 10, 64)
+	return index
 }