github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/sstable/compression.go (about)

     1  // Copyright 2021 The LevelDB-Go and Pebble Authors. All rights reserved. Use
     2  // of this source code is governed by a BSD-style license that can be found in
     3  // the LICENSE file.
     4  
     5  package sstable
     6  
     7  import (
     8  	"encoding/binary"
     9  
    10  	"github.com/cockroachdb/errors"
    11  	"github.com/cockroachdb/pebble/internal/base"
    12  	"github.com/cockroachdb/pebble/internal/cache"
    13  	"github.com/golang/snappy"
    14  )
    15  
    16  func decompressedLen(blockType blockType, b []byte) (int, int, error) {
    17  	switch blockType {
    18  	case noCompressionBlockType:
    19  		return 0, 0, nil
    20  	case snappyCompressionBlockType:
    21  		l, err := snappy.DecodedLen(b)
    22  		return l, 0, err
    23  	case zstdCompressionBlockType:
    24  		// This will also be used by zlib, bzip2 and lz4 to retrieve the decodedLen
    25  		// if we implement these algorithms in the future.
    26  		decodedLenU64, varIntLen := binary.Uvarint(b)
    27  		if varIntLen <= 0 {
    28  			return 0, 0, base.CorruptionErrorf("pebble/table: compression block has invalid length")
    29  		}
    30  		return int(decodedLenU64), varIntLen, nil
    31  	default:
    32  		return 0, 0, base.CorruptionErrorf("pebble/table: unknown block compression: %d", errors.Safe(blockType))
    33  	}
    34  }
    35  
    36  func decompressInto(blockType blockType, compressed []byte, buf []byte) ([]byte, error) {
    37  	var result []byte
    38  	var err error
    39  	switch blockType {
    40  	case snappyCompressionBlockType:
    41  		result, err = snappy.Decode(buf, compressed)
    42  	case zstdCompressionBlockType:
    43  		result, err = decodeZstd(buf, compressed)
    44  	}
    45  	if err != nil {
    46  		return nil, base.MarkCorruptionError(err)
    47  	}
    48  	if len(result) != 0 && (len(result) != len(buf) || &result[0] != &buf[0]) {
    49  		return nil, base.CorruptionErrorf("pebble/table: decompressed into unexpected buffer: %p != %p",
    50  			errors.Safe(result), errors.Safe(buf))
    51  	}
    52  	return result, nil
    53  }
    54  
    55  // decompressBlock decompresses an SST block, with manually-allocated space.
    56  // NB: If decompressBlock returns (nil, nil), no decompression was necessary and
    57  // the caller may use `b` directly.
    58  func decompressBlock(blockType blockType, b []byte) (*cache.Value, error) {
    59  	if blockType == noCompressionBlockType {
    60  		return nil, nil
    61  	}
    62  	// first obtain the decoded length.
    63  	decodedLen, prefixLen, err := decompressedLen(blockType, b)
    64  	if err != nil {
    65  		return nil, err
    66  	}
    67  	b = b[prefixLen:]
    68  	// Allocate sufficient space from the cache.
    69  	decoded := cache.Alloc(decodedLen)
    70  	decodedBuf := decoded.Buf()
    71  	if _, err := decompressInto(blockType, b, decodedBuf); err != nil {
    72  		cache.Free(decoded)
    73  		return nil, err
    74  	}
    75  	return decoded, nil
    76  }
    77  
    78  // compressBlock compresses an SST block, using compressBuf as the desired destination.
    79  func compressBlock(
    80  	compression Compression, b []byte, compressedBuf []byte,
    81  ) (blockType blockType, compressed []byte) {
    82  	switch compression {
    83  	case SnappyCompression:
    84  		return snappyCompressionBlockType, snappy.Encode(compressedBuf, b)
    85  	case NoCompression:
    86  		return noCompressionBlockType, b
    87  	}
    88  
    89  	if len(compressedBuf) < binary.MaxVarintLen64 {
    90  		compressedBuf = append(compressedBuf, make([]byte, binary.MaxVarintLen64-len(compressedBuf))...)
    91  	}
    92  	varIntLen := binary.PutUvarint(compressedBuf, uint64(len(b)))
    93  	switch compression {
    94  	case ZstdCompression:
    95  		return zstdCompressionBlockType, encodeZstd(compressedBuf, varIntLen, b)
    96  	default:
    97  		return noCompressionBlockType, b
    98  	}
    99  }