Merge pull request grafana#1044 from cortexproject/shrunk-chunks

 Write only as many chunk bytes as needed

chunk_test.go

@@ -16,6 +16,10 @@ import (
 
 const userID = "userID"
 
+func init() {
+	encoding.DefaultEncoding = encoding.Varbit
+}
+
 func dummyChunk(now model.Time) Chunk {
 	return dummyChunkFor(now, model.Metric{
 		model.MetricNameLabel: "foo",

encoding/bigchunk.go

@@ -5,7 +5,6 @@ import (
 	"encoding/binary"
 	"errors"
 	"io"
-	"io/ioutil"
 
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/tsdb/chunkenc"
@@ -95,14 +94,6 @@ func (b *bigchunk) MarshalToBuf(buf []byte) error {
 	return b.Marshal(writer)
 }
 
-func (b *bigchunk) Unmarshal(r io.Reader) error {
-	buf, err := ioutil.ReadAll(r)
-	if err != nil {
-		return err
-	}
-	return b.UnmarshalFromBuf(buf)
-}
-
 func (b *bigchunk) UnmarshalFromBuf(buf []byte) error {
 	r := reader{buf: buf}
 	numChunks, err := r.ReadUint16()

encoding/chunk.go

@@ -86,7 +86,6 @@ type Chunk interface {
 	Add(sample model.SamplePair) ([]Chunk, error)
 	NewIterator() Iterator
 	Marshal(io.Writer) error
-	Unmarshal(io.Reader) error
 	UnmarshalFromBuf([]byte) error
 	Encoding() Encoding
 	Utilization() float64

encoding/chunk_test.go

@@ -55,6 +55,7 @@ func TestLen(t *testing.T) {
 var step = int(15 * time.Second / time.Millisecond)
 
 func TestChunk(t *testing.T) {
+	alwaysMarshalFullsizeChunks = false
 	for _, tc := range []struct {
 		encoding   Encoding
 		maxSamples int
@@ -112,7 +113,7 @@ func testChunkEncoding(t *testing.T, encoding Encoding, samples int) {
 
 	bs1 := buf.Bytes()
 	chunk, err = NewForEncoding(encoding)
-	err = chunk.Unmarshal(&buf)
+	err = chunk.UnmarshalFromBuf(bs1)
 	require.NoError(t, err)
 
 	// Check all the samples are in there.
@@ -127,11 +128,12 @@ func testChunkEncoding(t *testing.T, encoding Encoding, samples int) {
 	require.NoError(t, iter.Err())
 
 	// Check the byte representation after another Marshall is the same.
+	buf = bytes.Buffer{}
 	err = chunk.Marshal(&buf)
 	require.NoError(t, err)
 	bs2 := buf.Bytes()
 
-	require.True(t, bytes.Equal(bs1, bs2))
+	require.Equal(t, bs1, bs2)
 }
 
 // testChunkSeek checks seek works as expected.

encoding/delta.go

@@ -213,15 +213,6 @@ func (c deltaEncodedChunk) Marshal(w io.Writer) error {
 	return nil
 }
 
-// Unmarshal implements chunk.
-func (c *deltaEncodedChunk) Unmarshal(r io.Reader) error {
-	*c = (*c)[:cap(*c)]
-	if _, err := io.ReadFull(r, *c); err != nil {
-		return err
-	}
-	return c.setLen()
-}
-
 // UnmarshalFromBuf implements chunk.
 func (c *deltaEncodedChunk) UnmarshalFromBuf(buf []byte) error {
 	*c = (*c)[:cap(*c)]

encoding/delta_test.go

@@ -99,9 +99,6 @@ func TestUnmarshalingCorruptedDeltaReturnsAnError(t *testing.T) {
 		err = cs[0].UnmarshalFromBuf(buf)
 		verifyUnmarshallingError(err, c.chunkTypeName, "buf", "invalid number of time bytes")
 
-		err = cs[0].Unmarshal(bytes.NewBuffer(buf))
-		verifyUnmarshallingError(err, c.chunkTypeName, "Reader", "invalid number of time bytes")
-
 		// Fix the corruption to go on.
 		buf[c.timeBytesPos] = byte(d1)
 
@@ -111,9 +108,6 @@ func TestUnmarshalingCorruptedDeltaReturnsAnError(t *testing.T) {
 
 			err = cs[0].UnmarshalFromBuf(buf)
 			verifyUnmarshallingError(err, c.chunkTypeName, "buf", "header size")
-
-			err = cs[0].Unmarshal(bytes.NewBuffer(buf))
-			verifyUnmarshallingError(err, c.chunkTypeName, "Reader", "header size")
 		}
 	}
 }

encoding/doubledelta.go

@@ -241,15 +241,6 @@ func (c doubleDeltaEncodedChunk) MarshalToBuf(buf []byte) error {
 	return nil
 }
 
-// Unmarshal implements chunk.
-func (c *doubleDeltaEncodedChunk) Unmarshal(r io.Reader) error {
-	*c = (*c)[:cap(*c)]
-	if _, err := io.ReadFull(r, *c); err != nil {
-		return err
-	}
-	return c.setLen()
-}
-
 // UnmarshalFromBuf implements chunk.
 func (c *doubleDeltaEncodedChunk) UnmarshalFromBuf(buf []byte) error {
 	*c = (*c)[:cap(*c)]

encoding/varbit.go

@@ -18,6 +18,7 @@ package encoding
 
 import (
 	"encoding/binary"
+	"flag"
 	"fmt"
 	"io"
 	"math"
@@ -286,26 +287,21 @@ func (c *varbitChunk) Slice(_, _ model.Time) Chunk {
 
 // Marshal implements chunk.
 func (c varbitChunk) Marshal(w io.Writer) error {
-	n, err := w.Write(c)
+	size := c.marshalLen()
+	n, err := w.Write(c[:size])
 	if err != nil {
 		return err
 	}
-	if n != cap(c) {
-		return fmt.Errorf("wanted to write %d bytes, wrote %d", cap(c), n)
+	if n != size {
+		return fmt.Errorf("wanted to write %d bytes, wrote %d", size, n)
 	}
 	return nil
 }
 
-// Unmarshal implements chunk.
-func (c varbitChunk) Unmarshal(r io.Reader) error {
-	_, err := io.ReadFull(r, c)
-	return err
-}
-
 // UnmarshalFromBuf implements chunk.
 func (c varbitChunk) UnmarshalFromBuf(buf []byte) error {
-	if copied := copy(c, buf); copied != cap(c) {
-		return fmt.Errorf("insufficient bytes copied from buffer during unmarshaling, want %d, got %d", cap(c), copied)
+	if copied := copy(c, buf); copied != cap(c) && copied != c.marshalLen() {
+		return fmt.Errorf("incorrect byte count copied from buffer during unmarshaling, want %d or %d, got %d", c.marshalLen(), ChunkLen, copied)
 	}
 	return nil
 }
@@ -319,6 +315,32 @@ func (c varbitChunk) Utilization() float64 {
 	return math.Min(float64(c.nextSampleOffset()/8+15)/float64(cap(c)), 1)
 }
 
+// MarshalConfig configures the behaviour of marshalling
+type MarshalConfig struct{}
+
+var alwaysMarshalFullsizeChunks = true
+
+// RegisterFlags registers configuration settings.
+func (MarshalConfig) RegisterFlags(f *flag.FlagSet) {
+	flag.BoolVar(&alwaysMarshalFullsizeChunks, "store.fullsize-chunks", alwaysMarshalFullsizeChunks, "When saving varbit chunks, pad to 1024 bytes")
+}
+
+// marshalLen returns the number of bytes that should be marshalled for this chunk
+func (c varbitChunk) marshalLen() int {
+	if alwaysMarshalFullsizeChunks {
+		return cap(c)
+	}
+	bits := c.nextSampleOffset()
+	if bits < varbitThirdSampleBitOffset {
+		bits = varbitThirdSampleBitOffset
+	}
+	bytes := int(bits)/8 + 1
+	if bytes > len(c) {
+		bytes = len(c)
+	}
+	return bytes
+}
+
 // Len implements chunk.  Runs in O(n).
 func (c varbitChunk) Len() int {
 	it := c.NewIterator()
@@ -329,7 +351,7 @@ func (c varbitChunk) Len() int {
 }
 
 func (c varbitChunk) Size() int {
-	return len(c)
+	return c.marshalLen()
 }
 
 func (c varbitChunk) firstTime() model.Time {