163 lines · 3.5 KB · Go
RawDownload
// 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(),	)}