summaryrefslogtreecommitdiffstats
path: root/modules/zstd/zstd.go
diff options
context:
space:
mode:
authorDaniel Baumann <daniel@debian.org>2024-10-18 20:33:49 +0200
committerDaniel Baumann <daniel@debian.org>2024-12-12 23:57:56 +0100
commite68b9d00a6e05b3a941f63ffb696f91e554ac5ec (patch)
tree97775d6c13b0f416af55314eb6a89ef792474615 /modules/zstd/zstd.go
parentInitial commit. (diff)
downloadforgejo-e68b9d00a6e05b3a941f63ffb696f91e554ac5ec.tar.xz
forgejo-e68b9d00a6e05b3a941f63ffb696f91e554ac5ec.zip
Adding upstream version 9.0.3.
Signed-off-by: Daniel Baumann <daniel@debian.org>
Diffstat (limited to 'modules/zstd/zstd.go')
-rw-r--r--modules/zstd/zstd.go163
1 files changed, 163 insertions, 0 deletions
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(),
+ )
+}