github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/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/golang/snappy"
    12  	"github.com/zuoyebang/bitalostable/internal/base"
    13  	"github.com/zuoyebang/bitalostable/internal/cache"
    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("bitalostable/table: compression block has invalid length")
    29  		}
    30  		return int(decodedLenU64), varIntLen, nil
    31  	default:
    32  		return 0, 0, base.CorruptionErrorf("bitalostable/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("bitalostable/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 space allocated from a cache.
    56  func decompressBlock(cache *cache.Cache, blockType blockType, b []byte) (*cache.Value, error) {
    57  	if blockType == noCompressionBlockType {
    58  		return nil, nil
    59  	}
    60  	// first obtain the decoded length.
    61  	decodedLen, prefixLen, err := decompressedLen(blockType, b)
    62  	if err != nil {
    63  		return nil, err
    64  	}
    65  	if prefixLen != 0 {
    66  		b = b[prefixLen:]
    67  	}
    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  	}
    74  	return decoded, nil
    75  }
    76  
    77  // compressBlock compresses an SST block, using compressBuf as the desired destination.
    78  func compressBlock(
    79  	compression Compression, b []byte, compressedBuf []byte,
    80  ) (blockType blockType, compressed []byte) {
    81  	switch compression {
    82  	case SnappyCompression:
    83  		return snappyCompressionBlockType, snappy.Encode(compressedBuf, b)
    84  	case NoCompression:
    85  		return noCompressionBlockType, b
    86  	}
    87  
    88  	if len(compressedBuf) < binary.MaxVarintLen64 {
    89  		compressedBuf = append(compressedBuf, make([]byte, binary.MaxVarintLen64-len(compressedBuf))...)
    90  	}
    91  	varIntLen := binary.PutUvarint(compressedBuf, uint64(len(b)))
    92  	switch compression {
    93  	case ZstdCompression:
    94  		return zstdCompressionBlockType, encodeZstd(compressedBuf, varIntLen, b)
    95  	default:
    96  		return noCompressionBlockType, b
    97  	}
    98  }