github.com/thanos-io/thanos@v0.32.5/pkg/compact/downsample/pool.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package downsample
     5  
     6  import (
     7  	"sync"
     8  
     9  	"github.com/prometheus/prometheus/tsdb/chunkenc"
    10  )
    11  
    12  // Pool is a memory pool of chunk objects, supporting Thanos aggregated chunk encoding.
    13  // It maintains separate pools for xor and aggr chunks.
    14  type pool struct {
    15  	wrapped chunkenc.Pool
    16  	aggr    sync.Pool
    17  }
    18  
    19  // TODO(bwplotka): Add reasonable limits to our sync pooling them to detect OOMs early.
    20  func NewPool() chunkenc.Pool {
    21  	return &pool{
    22  		wrapped: chunkenc.NewPool(),
    23  		aggr: sync.Pool{
    24  			New: func() interface{} {
    25  				return &AggrChunk{}
    26  			},
    27  		},
    28  	}
    29  }
    30  
    31  func (p *pool) Get(e chunkenc.Encoding, b []byte) (chunkenc.Chunk, error) {
    32  	switch e {
    33  	case ChunkEncAggr:
    34  		c := p.aggr.Get().(*AggrChunk)
    35  		*c = AggrChunk(b)
    36  		return c, nil
    37  	}
    38  
    39  	return p.wrapped.Get(e, b)
    40  
    41  }
    42  
    43  func (p *pool) Put(c chunkenc.Chunk) error {
    44  	switch c.Encoding() {
    45  	case ChunkEncAggr:
    46  		ac, ok := c.(*AggrChunk)
    47  		// This may happen often with wrapped chunks. Nothing we can really do about
    48  		// it but returning an error would cause a lot of allocations again. Thus,
    49  		// we just skip it.
    50  		if !ok {
    51  			return nil
    52  		}
    53  
    54  		// Clear []byte.
    55  		*ac = AggrChunk(nil)
    56  		p.aggr.Put(ac)
    57  		return nil
    58  	}
    59  
    60  	return p.wrapped.Put(c)
    61  }