128 lines
		
	
	
		
			3.9 KiB
		
	
	
	
		
			Go
		
	
	
	
			
		
		
	
	
			128 lines
		
	
	
		
			3.9 KiB
		
	
	
	
		
			Go
		
	
	
	
| // Copyright 2020 The Prometheus Authors
 | |
| // Licensed under the Apache License, Version 2.0 (the "License");
 | |
| // you may not use this file except in compliance with the License.
 | |
| // You may obtain a copy of the License at
 | |
| //
 | |
| // http://www.apache.org/licenses/LICENSE-2.0
 | |
| //
 | |
| // Unless required by applicable law or agreed to in writing, software
 | |
| // distributed under the License is distributed on an "AS IS" BASIS,
 | |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 | |
| // See the License for the specific language governing permissions and
 | |
| // limitations under the License.
 | |
| 
 | |
| package tsdb
 | |
| 
 | |
| import (
 | |
| 	"context"
 | |
| 	"errors"
 | |
| 	"fmt"
 | |
| 	"log/slog"
 | |
| 	"math"
 | |
| 	"os"
 | |
| 
 | |
| 	"github.com/oklog/ulid/v2"
 | |
| 
 | |
| 	"github.com/prometheus/prometheus/model/timestamp"
 | |
| 	"github.com/prometheus/prometheus/storage"
 | |
| 	"github.com/prometheus/prometheus/tsdb/chunkenc"
 | |
| )
 | |
| 
 | |
| // BlockWriter is a block writer that allows appending and flushing series to disk.
 | |
| type BlockWriter struct {
 | |
| 	logger         *slog.Logger
 | |
| 	destinationDir string
 | |
| 
 | |
| 	head      *Head
 | |
| 	blockSize int64 // in ms
 | |
| 	chunkDir  string
 | |
| }
 | |
| 
 | |
| // ErrNoSeriesAppended is returned if the series count is zero while flushing blocks.
 | |
| var ErrNoSeriesAppended = errors.New("no series appended, aborting")
 | |
| 
 | |
| // NewBlockWriter creates a new block writer.
 | |
| //
 | |
| // The returned writer accumulates all the series in the Head block until `Flush` is called.
 | |
| //
 | |
| // Note that the writer will not check if the target directory exists or
 | |
| // contains anything at all. It is the caller's responsibility to
 | |
| // ensure that the resulting blocks do not overlap etc.
 | |
| // Writer ensures the block flush is atomic (via rename).
 | |
| func NewBlockWriter(logger *slog.Logger, dir string, blockSize int64) (*BlockWriter, error) {
 | |
| 	w := &BlockWriter{
 | |
| 		logger:         logger,
 | |
| 		destinationDir: dir,
 | |
| 		blockSize:      blockSize,
 | |
| 	}
 | |
| 	if err := w.initHead(); err != nil {
 | |
| 		return nil, err
 | |
| 	}
 | |
| 	return w, nil
 | |
| }
 | |
| 
 | |
| // initHead creates and initialises a new TSDB head.
 | |
| func (w *BlockWriter) initHead() error {
 | |
| 	chunkDir, err := os.MkdirTemp(os.TempDir(), "head")
 | |
| 	if err != nil {
 | |
| 		return fmt.Errorf("create temp dir: %w", err)
 | |
| 	}
 | |
| 	w.chunkDir = chunkDir
 | |
| 	opts := DefaultHeadOptions()
 | |
| 	opts.ChunkRange = w.blockSize
 | |
| 	opts.ChunkDirRoot = w.chunkDir
 | |
| 	opts.EnableNativeHistograms.Store(true)
 | |
| 	h, err := NewHead(nil, w.logger, nil, nil, opts, NewHeadStats())
 | |
| 	if err != nil {
 | |
| 		return fmt.Errorf("tsdb.NewHead: %w", err)
 | |
| 	}
 | |
| 
 | |
| 	w.head = h
 | |
| 	return w.head.Init(math.MinInt64)
 | |
| }
 | |
| 
 | |
| // Appender returns a new appender on the database.
 | |
| // Appender can't be called concurrently. However, the returned Appender can safely be used concurrently.
 | |
| func (w *BlockWriter) Appender(ctx context.Context) storage.Appender {
 | |
| 	return w.head.Appender(ctx)
 | |
| }
 | |
| 
 | |
| // Flush implements the Writer interface. This is where actual block writing
 | |
| // happens. After flush completes, no writes can be done.
 | |
| func (w *BlockWriter) Flush(ctx context.Context) (ulid.ULID, error) {
 | |
| 	mint := w.head.MinTime()
 | |
| 	// Add +1 millisecond to block maxt because block intervals are half-open: [b.MinTime, b.MaxTime).
 | |
| 	// Because of this block intervals are always +1 than the total samples it includes.
 | |
| 	maxt := w.head.MaxTime() + 1
 | |
| 	w.logger.Info("flushing", "series_count", w.head.NumSeries(), "mint", timestamp.Time(mint), "maxt", timestamp.Time(maxt))
 | |
| 
 | |
| 	compactor, err := NewLeveledCompactor(ctx,
 | |
| 		nil,
 | |
| 		w.logger,
 | |
| 		[]int64{w.blockSize},
 | |
| 		chunkenc.NewPool(), nil)
 | |
| 	if err != nil {
 | |
| 		return ulid.ULID{}, fmt.Errorf("create leveled compactor: %w", err)
 | |
| 	}
 | |
| 	ids, err := compactor.Write(w.destinationDir, w.head, mint, maxt, nil)
 | |
| 	if err != nil {
 | |
| 		return ulid.ULID{}, fmt.Errorf("compactor write: %w", err)
 | |
| 	}
 | |
| 
 | |
| 	// No block was produced. Caller is responsible to check empty
 | |
| 	// ulid.ULID based on its use case.
 | |
| 	if len(ids) == 0 {
 | |
| 		return ulid.ULID{}, nil
 | |
| 	}
 | |
| 	return ids[0], nil
 | |
| }
 | |
| 
 | |
| func (w *BlockWriter) Close() error {
 | |
| 	defer func() {
 | |
| 		if err := os.RemoveAll(w.chunkDir); err != nil {
 | |
| 			w.logger.Error("error in deleting BlockWriter files", "err", err)
 | |
| 		}
 | |
| 	}()
 | |
| 	return w.head.Close()
 | |
| }
 |