Atlas/internal/store/write_queue.go

package store

import (
	"context"
	"log"
	"sync"
	"time"
)

type writeJob func(ctx context.Context, db *Store) error

type WriteQueue struct {
	store    *Store
	ch       chan writeJob
	wg       sync.WaitGroup
	interval time.Duration
	stopCh   chan struct{}
}

func NewWriteQueue(s *Store, flushInterval time.Duration) *WriteQueue {
	wq := &WriteQueue{
		store:    s,
		ch:       make(chan writeJob, 256),
		interval: flushInterval,
		stopCh:   make(chan struct{}),
	}
	wq.wg.Add(1)
	go wq.run()
	return wq
}

func (wq *WriteQueue) Enqueue(job writeJob) {
	select {
	case wq.ch <- job:
	default:
		go func() { wq.ch <- job }()
	}
}

func (wq *WriteQueue) ForceFlush(ctx context.Context, job writeJob) error {
	return job(ctx, wq.store)
}

func (wq *WriteQueue) run() {
	defer wq.wg.Done()
	ticker := time.NewTicker(wq.interval)
	defer ticker.Stop()
	pending := make([]writeJob, 0, 32)
	flush := func() {
		if len(pending) == 0 {
			return
		}
		jobs := pending
		pending = pending[:0]
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		for _, job := range jobs {
			if err := job(ctx, wq.store); err != nil {
				log.Printf("write queue: %v", err)
			}
		}
	}
	for {
		select {
		case <-wq.stopCh:
			flush()
			for {
				select {
				case job := <-wq.ch:
					pending = append(pending, job)
				default:
					flush()
					return
				}
			}
		case job := <-wq.ch:
			pending = append(pending, job)
			if len(pending) >= 32 {
				flush()
			}
		case <-ticker.C:
			flush()
		}
	}
}

func (wq *WriteQueue) Close() {
	close(wq.stopCh)
	wq.wg.Wait()
}