go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/filter/dscache/serialize.go (about)

     1  // Copyright 2015 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package dscache
    16  
    17  import (
    18  	"bytes"
    19  	"compress/zlib"
    20  	"io"
    21  
    22  	"go.chromium.org/luci/common/errors"
    23  
    24  	"go.chromium.org/luci/gae/internal/zstd"
    25  	ds "go.chromium.org/luci/gae/service/datastore"
    26  )
    27  
    28  // UseZstd, if true, instructs the dscache to use zstd algorithm for compression
    29  // instead of zlib.
    30  //
    31  // Already compressed zlib entities are supported even when UseZstd is true.
    32  var UseZstd = false
    33  
    34  type compressionType byte
    35  
    36  const (
    37  	compressionNone compressionType = 0
    38  	compressionZlib compressionType = 1
    39  	compressionZstd compressionType = 2
    40  )
    41  
    42  func encodeItemValue(pm ds.PropertyMap, pfx []byte) ([]byte, error) {
    43  	var err error
    44  	if pm, err = pm.Save(false); err != nil {
    45  		return nil, err
    46  	}
    47  
    48  	// Try to write as uncompressed first. Capacity of 256 is picked arbitrarily.
    49  	// Most entities are pretty small.
    50  	buf := bytes.NewBuffer(make([]byte, 0, 256))
    51  	buf.Write(pfx)
    52  	buf.WriteByte(byte(compressionNone))
    53  	if err := ds.Serialize.PropertyMap(buf, pm); err != nil {
    54  		return nil, err
    55  	}
    56  
    57  	// If it is small enough, we are done.
    58  	if buf.Len() <= CompressionThreshold {
    59  		return buf.Bytes(), nil
    60  	}
    61  
    62  	// If too big, grab a new buffer and compress data there. Preallocate a new
    63  	// buffer assuming 2x compression.
    64  	data := buf.Bytes()[len(pfx)+1:] // skip pfx and compressionNone byte
    65  	buf2 := bytes.NewBuffer(make([]byte, 0, len(data)/2))
    66  
    67  	// Use zlib in legacy applications that didn't opt-in into using zstd.
    68  	algo := compressionZlib
    69  	if UseZstd {
    70  		algo = compressionZstd
    71  	}
    72  
    73  	// Compress into the new buffer.
    74  	buf2.Write(pfx)
    75  	buf2.WriteByte(byte(algo))
    76  	switch algo {
    77  	case compressionZlib:
    78  		writer := zlib.NewWriter(buf2)
    79  		writer.Write(data)
    80  		writer.Close()
    81  		return buf2.Bytes(), nil
    82  	case compressionZstd:
    83  		return zstd.EncodeAll(data, buf2.Bytes()), nil
    84  	default:
    85  		panic("impossible")
    86  	}
    87  }
    88  
    89  func decodeItemValue(val []byte, kc ds.KeyContext) (ds.PropertyMap, error) {
    90  	if len(val) == 0 {
    91  		return nil, ds.ErrNoSuchEntity
    92  	}
    93  	if len(val) < 1 {
    94  		return nil, errors.Reason("missing the compression type byte").Err()
    95  	}
    96  
    97  	compTypeByte, data := val[0], val[1:]
    98  
    99  	switch compressionType(compTypeByte) {
   100  	case compressionNone:
   101  		// already decompressed
   102  
   103  	case compressionZlib:
   104  		reader, err := zlib.NewReader(bytes.NewBuffer(data))
   105  		if err != nil {
   106  			return nil, err
   107  		}
   108  		if data, err = io.ReadAll(reader); err != nil {
   109  			return nil, err
   110  		}
   111  		if err = reader.Close(); err != nil {
   112  			return nil, err
   113  		}
   114  
   115  	case compressionZstd:
   116  		var err error
   117  		data, err = zstd.DecodeAll(data, nil)
   118  		if err != nil {
   119  			return nil, err
   120  		}
   121  
   122  	default:
   123  		return nil, errors.Reason("unknown compression scheme #%d", compTypeByte).Err()
   124  	}
   125  
   126  	return ds.Deserializer{KeyContext: kc}.PropertyMap(bytes.NewBuffer(data))
   127  }