From dd136858f1ea40ad3c94191d647487fa4f31926c Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Fri, 18 Oct 2024 20:33:49 +0200 Subject: Adding upstream version 9.0.0. Signed-off-by: Daniel Baumann --- modules/zstd/zstd.go | 163 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 163 insertions(+) create mode 100644 modules/zstd/zstd.go (limited to 'modules/zstd/zstd.go') diff --git a/modules/zstd/zstd.go b/modules/zstd/zstd.go new file mode 100644 index 0000000..d224944 --- /dev/null +++ b/modules/zstd/zstd.go @@ -0,0 +1,163 @@ +// Copyright 2024 The Gitea Authors. All rights reserved. +// SPDX-License-Identifier: MIT + +// Package zstd provides a high-level API for reading and writing zstd-compressed data. +// It supports both regular and seekable zstd streams. +// It's not a new wheel, but a wrapper around the zstd and zstd-seekable-format-go packages. +package zstd + +import ( + "errors" + "io" + + seekable "github.com/SaveTheRbtz/zstd-seekable-format-go/pkg" + "github.com/klauspost/compress/zstd" +) + +type Writer zstd.Encoder + +var _ io.WriteCloser = (*Writer)(nil) + +// NewWriter returns a new zstd writer. +func NewWriter(w io.Writer, opts ...WriterOption) (*Writer, error) { + zstdW, err := zstd.NewWriter(w, opts...) + if err != nil { + return nil, err + } + return (*Writer)(zstdW), nil +} + +func (w *Writer) Write(p []byte) (int, error) { + return (*zstd.Encoder)(w).Write(p) +} + +func (w *Writer) Close() error { + return (*zstd.Encoder)(w).Close() +} + +type Reader zstd.Decoder + +var _ io.ReadCloser = (*Reader)(nil) + +// NewReader returns a new zstd reader. +func NewReader(r io.Reader, opts ...ReaderOption) (*Reader, error) { + zstdR, err := zstd.NewReader(r, opts...) + if err != nil { + return nil, err + } + return (*Reader)(zstdR), nil +} + +func (r *Reader) Read(p []byte) (int, error) { + return (*zstd.Decoder)(r).Read(p) +} + +func (r *Reader) Close() error { + (*zstd.Decoder)(r).Close() // no error returned + return nil +} + +type SeekableWriter struct { + buf []byte + n int + w seekable.Writer +} + +var _ io.WriteCloser = (*SeekableWriter)(nil) + +// NewSeekableWriter returns a zstd writer to compress data to seekable format. +// blockSize is an important parameter, it should be decided according to the actual business requirements. +// If it's too small, the compression ratio could be very bad, even no compression at all. +// If it's too large, it could cost more traffic when reading the data partially from underlying storage. +func NewSeekableWriter(w io.Writer, blockSize int, opts ...WriterOption) (*SeekableWriter, error) { + zstdW, err := zstd.NewWriter(nil, opts...) + if err != nil { + return nil, err + } + + seekableW, err := seekable.NewWriter(w, zstdW) + if err != nil { + return nil, err + } + + return &SeekableWriter{ + buf: make([]byte, blockSize), + w: seekableW, + }, nil +} + +func (w *SeekableWriter) Write(p []byte) (int, error) { + written := 0 + for len(p) > 0 { + n := copy(w.buf[w.n:], p) + w.n += n + written += n + p = p[n:] + + if w.n == len(w.buf) { + if _, err := w.w.Write(w.buf); err != nil { + return written, err + } + w.n = 0 + } + } + return written, nil +} + +func (w *SeekableWriter) Close() error { + if w.n > 0 { + if _, err := w.w.Write(w.buf[:w.n]); err != nil { + return err + } + } + return w.w.Close() +} + +type SeekableReader struct { + r seekable.Reader + c func() error +} + +var _ io.ReadSeekCloser = (*SeekableReader)(nil) + +// NewSeekableReader returns a zstd reader to decompress data from seekable format. +func NewSeekableReader(r io.ReadSeeker, opts ...ReaderOption) (*SeekableReader, error) { + zstdR, err := zstd.NewReader(nil, opts...) + if err != nil { + return nil, err + } + + seekableR, err := seekable.NewReader(r, zstdR) + if err != nil { + return nil, err + } + + ret := &SeekableReader{ + r: seekableR, + } + if closer, ok := r.(io.Closer); ok { + ret.c = closer.Close + } + + return ret, nil +} + +func (r *SeekableReader) Read(p []byte) (int, error) { + return r.r.Read(p) +} + +func (r *SeekableReader) Seek(offset int64, whence int) (int64, error) { + return r.r.Seek(offset, whence) +} + +func (r *SeekableReader) Close() error { + return errors.Join( + func() error { + if r.c != nil { + return r.c() + } + return nil + }(), + r.r.Close(), + ) +} -- cgit v1.2.3