191 lines · 5.2 KB · Go
RawDownload
package compress import (	"bufio"	"context"	"crypto/md5"	"encoding/hex"	"errors"	"io" 	"github.com/klauspost/compress/zstd" 	"github.com/rclone/rclone/fs"	"github.com/rclone/rclone/fs/accounting"	"github.com/rclone/rclone/fs/chunkedreader"	"github.com/rclone/rclone/fs/hash") // zstdModeHandler implements compressionModeHandler for zstdtype zstdModeHandler struct{} // isCompressible checks the compression ratio of the provided data and returns true if the ratio exceeds// the configured thresholdfunc (z *zstdModeHandler) isCompressible(r io.Reader, compressionMode int) (bool, error) {	var b countingDiscard	var n int64	w, err := NewWriterSzstd(&b, zstd.WithEncoderLevel(zstd.SpeedDefault))	if err != nil {		return false, err	}	n, err = io.Copy(w, r)	if err != nil {		return false, err	}	err = w.Close()	if err != nil {		return false, err	}	ratio := float64(n) / float64(b)	return ratio > minCompressionRatio, nil} // newObjectGetOriginalSize returns the original file size from the metadatafunc (z *zstdModeHandler) newObjectGetOriginalSize(meta *ObjectMetadata) (int64, error) {	if meta.CompressionMetadataZstd == nil {		return 0, errors.New("missing zstd metadata")	}	return meta.CompressionMetadataZstd.Size, nil} // openGetReadCloser opens a compressed object and returns a ReadCloser in the Open methodfunc (z *zstdModeHandler) openGetReadCloser(	ctx context.Context,	o *Object,	offset int64,	limit int64,	cr chunkedreader.ChunkedReader,	closer io.Closer,	options ...fs.OpenOption,) (rc io.ReadCloser, err error) {	var file io.Reader 	if offset != 0 {		file, err = NewReaderAtSzstd(cr, o.meta.CompressionMetadataZstd, offset)	} else {		file, err = zstd.NewReader(cr)	}	if err != nil {		return nil, err	} 	var fileReader io.Reader	if limit != -1 {		fileReader = io.LimitReader(file, limit)	} else {		fileReader = file	}	// Return a ReadCloser	return ReadCloserWrapper{Reader: fileReader, Closer: closer}, nil} // processFileNameGetFileExtension returns the file extension for the given compression modefunc (z *zstdModeHandler) processFileNameGetFileExtension(compressionMode int) string {	if compressionMode == Zstd {		return zstdFileExt	} 	return ""} // putCompress compresses the input data and uploads it to the remote, returning the new object and its metadatafunc (z *zstdModeHandler) putCompress(	ctx context.Context,	f *Fs,	in io.Reader,	src fs.ObjectInfo,	options []fs.OpenOption,	mimeType string,) (fs.Object, *ObjectMetadata, error) {	// Unwrap reader accounting	in, wrap := accounting.UnWrap(in) 	// Add the metadata hasher	metaHasher := md5.New()	in = io.TeeReader(in, metaHasher) 	// Compress the file	pipeReader, pipeWriter := io.Pipe() 	resultsZstd := make(chan compressionResult[SzstdMetadata])	go func() {		writer, err := NewWriterSzstd(pipeWriter, zstd.WithEncoderLevel(zstd.EncoderLevel(f.opt.CompressionLevel)))		if err != nil {			resultsZstd <- compressionResult[SzstdMetadata]{err: err}			close(resultsZstd)			return		}		_, err = io.Copy(writer, in)		if wErr := writer.Close(); wErr != nil && err == nil {			err = wErr		}		if cErr := pipeWriter.Close(); cErr != nil && err == nil {			err = cErr		} 		resultsZstd <- compressionResult[SzstdMetadata]{err: err, meta: writer.GetMetadata()}		close(resultsZstd)	}() 	wrappedIn := wrap(bufio.NewReaderSize(pipeReader, bufferSize)) 	ht := f.Fs.Hashes().GetOne()	var hasher *hash.MultiHasher	var err error	if ht != hash.None {		wrappedIn, wrap = accounting.UnWrap(wrappedIn)		hasher, err = hash.NewMultiHasherTypes(hash.NewHashSet(ht))		if err != nil {			return nil, nil, err		}		wrappedIn = io.TeeReader(wrappedIn, hasher)		wrappedIn = wrap(wrappedIn)	} 	o, err := f.rcat(ctx, makeDataName(src.Remote(), src.Size(), f.mode), io.NopCloser(wrappedIn), src.ModTime(ctx), options)	if err != nil {		return nil, nil, err	} 	result := <-resultsZstd	if result.err != nil {		if o != nil {			_ = o.Remove(ctx)		}		return nil, nil, result.err	} 	// Build metadata using uncompressed size for filename	meta := z.newMetadata(result.meta.Size, f.mode, result.meta, hex.EncodeToString(metaHasher.Sum(nil)), mimeType)	if ht != hash.None && hasher != nil {		err = f.verifyObjectHash(ctx, o, hasher, ht)		if err != nil {			return nil, nil, err		}	}	return o, meta, nil} // putUncompressGetNewMetadata returns metadata in the putUncompress method for a specific compression algorithmfunc (z *zstdModeHandler) putUncompressGetNewMetadata(o fs.Object, mode int, md5 string, mimeType string, sum []byte) (fs.Object, *ObjectMetadata, error) {	return o, z.newMetadata(o.Size(), mode, SzstdMetadata{}, hex.EncodeToString(sum), mimeType), nil} // This function generates a metadata object for sgzip.GzipMetadata or SzstdMetadata.// Warning: This function panics if cmeta is not of the expected type.func (z *zstdModeHandler) newMetadata(size int64, mode int, cmeta any, md5 string, mimeType string) *ObjectMetadata {	meta, ok := cmeta.(SzstdMetadata)	if !ok {		panic("invalid cmeta type: expected SzstdMetadata")	} 	objMeta := new(ObjectMetadata)	objMeta.Size = size	objMeta.Mode = mode	objMeta.CompressionMetadataGzip = nil	objMeta.CompressionMetadataZstd = &meta	objMeta.MD5 = md5	objMeta.MimeType = mimeType 	return objMeta}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t