Job queue

tsdb/chunks/chunk_write_queue.go

@@ -29,6 +29,10 @@ const (
 
 	// Minimum interval between shrinking of chunkWriteQueue.chunkRefMap.
 	chunkRefMapMinShrinkInterval = 10 * time.Minute
+
+	// Maximum size of segment used by job queue (number of elements). With chunkWriteJob being 64 bytes,
+	// this will use ~512 KiB for empty queue.
+	maxChunkQueueSegmentSize = 8192
 )
 
 type chunkWriteJob struct {
@@ -45,7 +49,7 @@ type chunkWriteJob struct {
 // Chunks that shall be written get added to the queue, which is consumed asynchronously.
 // Adding jobs to the queue is non-blocking as long as the queue isn't full.
 type chunkWriteQueue struct {
-	jobs chan chunkWriteJob
+	jobs *writeJobQueue
 
 	chunkRefMapMtx        sync.RWMutex
 	chunkRefMap           map[ChunkDiskMapperRef]chunkenc.Chunk
@@ -83,8 +87,13 @@ func newChunkWriteQueue(reg prometheus.Registerer, size int, writeChunk writeChu
 		[]string{"operation"},
 	)
 
+	segmentSize := size
+	if segmentSize > maxChunkQueueSegmentSize {
+		segmentSize = maxChunkQueueSegmentSize
+	}
+
 	q := &chunkWriteQueue{
-		jobs:                  make(chan chunkWriteJob, size),
+		jobs:                  newWriteJobQueue(size, segmentSize),
 		chunkRefMap:           make(map[ChunkDiskMapperRef]chunkenc.Chunk),
 		chunkRefMapLastShrink: time.Now(),
 		writeChunk:            writeChunk,
@@ -108,7 +117,12 @@ func (c *chunkWriteQueue) start() {
 	go func() {
 		defer c.workerWg.Done()
 
-		for job := range c.jobs {
+		for {
+			job, ok := c.jobs.pop()
+			if !ok {
+				return
+			}
+
 			c.processJob(job)
 		}
 	}()
@@ -191,7 +205,14 @@ func (c *chunkWriteQueue) addJob(job chunkWriteJob) (err error) {
 	}
 	c.chunkRefMapMtx.Unlock()
 
-	c.jobs <- job
+	ok := c.jobs.push(job)
+	if !ok {
+		c.chunkRefMapMtx.Lock()
+		delete(c.chunkRefMap, job.ref)
+		c.chunkRefMapMtx.Unlock()
+
+		return errors.New("queue is closed")
+	}
 
 	return nil
 }
@@ -218,7 +239,7 @@ func (c *chunkWriteQueue) stop() {
 
 	c.isRunning = false
 
-	close(c.jobs)
+	c.jobs.close()
 
 	c.workerWg.Wait()
 }
@@ -230,7 +251,7 @@ func (c *chunkWriteQueue) queueIsEmpty() bool {
 func (c *chunkWriteQueue) queueIsFull() bool {
 	// When the queue is full and blocked on the writer the chunkRefMap has one more job than the cap of the jobCh
 	// because one job is currently being processed and blocked in the writer.
-	return c.queueSize() == cap(c.jobs)+1
+	return c.queueSize() == c.jobs.maxSize+1
 }
 
 func (c *chunkWriteQueue) queueSize() int {

tsdb/chunks/queue.go

@@ -0,0 +1,127 @@
+package chunks
+
+import "sync"
+
+// writeJobQueue is similar to buffered channel of chunkWriteJob, but manages its own buffers
+// to avoid using a lot of memory when it's empty. It does that by storing elements into segments
+// of equal size (segmentSize). When segment is not used anymore, reference to it are removed,
+// so it can be treated as a garbage.
+type writeJobQueue struct {
+	maxSize     int
+	segmentSize int
+
+	mtx            sync.Mutex            // protects all following variables
+	pushed, popped *sync.Cond            // signalled when something is pushed into the queue or popped from it
+	first, last    *writeJobQueueSegment // pointer to first and last segment, if any
+	size           int                   // total size of the queue
+	closed         bool                  // after closing the queue, nothing can be pushed to it
+}
+
+type writeJobQueueSegment struct {
+	segment             []chunkWriteJob
+	nextRead, nextWrite int                   // index of next read and next write in this segment.
+	nextSegment         *writeJobQueueSegment // next segment, if any
+}
+
+func newWriteJobQueue(maxSize, segmentSize int) *writeJobQueue {
+	if maxSize <= 0 || segmentSize <= 0 {
+		panic("invalid queue")
+	}
+
+	q := &writeJobQueue{
+		maxSize:     maxSize,
+		segmentSize: segmentSize,
+	}
+
+	q.pushed = sync.NewCond(&q.mtx)
+	q.popped = sync.NewCond(&q.mtx)
+	return q
+}
+
+func (q *writeJobQueue) close() {
+	q.mtx.Lock()
+	defer q.mtx.Unlock()
+
+	q.closed = true
+
+	// unblock all blocked goroutines
+	q.pushed.Broadcast()
+	q.popped.Broadcast()
+}
+
+// push blocks until there is space available in the queue, and then adds job to the queue.
+// If queue is closed or gets closed while waiting for space, push returns false.
+func (q *writeJobQueue) push(job chunkWriteJob) bool {
+	q.mtx.Lock()
+	defer q.mtx.Unlock()
+
+	// wait until queue has more space or is closed
+	for !q.closed && q.size >= q.maxSize {
+		q.popped.Wait()
+	}
+
+	if q.closed {
+		return false
+	}
+
+	// Check if this segment has more space for writing, and create new one if not.
+	if q.last == nil || q.last.nextWrite >= q.segmentSize {
+		prevLast := q.last
+		q.last = &writeJobQueueSegment{
+			segment: make([]chunkWriteJob, q.segmentSize),
+		}
+
+		if prevLast != nil {
+			prevLast.nextSegment = q.last
+		}
+		if q.first == nil {
+			q.first = q.last
+		}
+	}
+
+	q.last.segment[q.last.nextWrite] = job
+	q.last.nextWrite++
+	q.size++
+	q.pushed.Signal()
+	return true
+}
+
+// pop returns first job from the queue, and true.
+// if queue is empty, pop blocks until there is a job (returns true), or until queue is closed (returns false).
+// If queue was already closed, pop first returns all remaining elements from the queue (with true value), and only then returns false.
+func (q *writeJobQueue) pop() (chunkWriteJob, bool) {
+	q.mtx.Lock()
+	defer q.mtx.Unlock()
+
+	// wait until something is pushed to the queue, or queue is closed.
+	for q.size == 0 {
+		if q.closed {
+			return chunkWriteJob{}, false
+		}
+
+		q.pushed.Wait()
+	}
+
+	res := q.first.segment[q.first.nextRead]
+	q.first.segment[q.first.nextRead] = chunkWriteJob{} // clear just-read element
+	q.first.nextRead++
+	q.size--
+
+	// If we have read all possible elements from first segment, we can drop it.
+	if q.first.nextRead >= q.segmentSize {
+		q.first = q.first.nextSegment
+		if q.first == nil {
+			q.last = nil
+		}
+	}
+
+	q.popped.Signal()
+	return res, true
+}
+
+func (q *writeJobQueue) length() int {
+	q.mtx.Lock()
+	defer q.mtx.Unlock()
+
+	return q.size
+}