Commit 120e1557 authored by Kamil Trzciński's avatar Kamil Trzciński
Browse files

Merge branch 'ajwalker/cache-perf/interface' into 'master'

Implement archiver/extractor interface

See merge request gitlab-org/gitlab-runner!2195
parents b2a5ff18 bef25c22
package archive
import (
"context"
"errors"
"fmt"
"io"
"os"
)
var (
// ErrUnsupportedArchiveFormat is returned if an archiver or extractor format
// requested has not been registered.
ErrUnsupportedArchiveFormat = errors.New("unsupported archive format")
)
// CompressionLevel type for specifying a compression level.
type CompressionLevel int
// Compression levels from fastest (low/zero compression ratio) to slowest
// (high compression ratio).
const (
FastestCompression CompressionLevel = -2
FastCompression CompressionLevel = -1
DefaultCompression CompressionLevel = 0
SlowCompression CompressionLevel = 1
SlowestCompression CompressionLevel = 2
)
// Format type for specifying format.
type Format string
// Formats supported by GitLab.
const (
Raw Format = "raw"
Gzip Format = "gzip"
Zip Format = "zip"
)
var (
archivers = make(map[Format]NewArchiverFunc)
extractors = make(map[Format]NewExtractorFunc)
)
// Archiver is an interface for the Archive method.
type Archiver interface {
Archive(ctx context.Context, files map[string]os.FileInfo) error
}
// Extractor is an interface for the Extract method.
type Extractor interface {
Extract(ctx context.Context) error
}
// NewArchiverFunc is a function that can be registered (with Register()) and
// used to instantiate a new archiver (with NewArchiver()).
type NewArchiverFunc func(w io.Writer, dir string, level CompressionLevel) (Archiver, error)
// NewExtractorFunc is a function that can be registered (with Register()) and
// used to instantiate a new extractor (with NewExtractor()).
type NewExtractorFunc func(r io.ReaderAt, size int64, dir string) (Extractor, error)
// Register registers a new archiver, overriding the archiver and/or extractor
// for the format provided.
func Register(format Format, archiver NewArchiverFunc, extractor NewExtractorFunc) {
if archiver != nil {
archivers[format] = archiver
}
if extractor != nil {
extractors[format] = extractor
}
}
// NewArchiver returns a new Archiver of the specified format.
//
// The archiver will ensure that files to be archived are children of the
// directory provided.
func NewArchiver(format Format, w io.Writer, dir string, level CompressionLevel) (Archiver, error) {
fn := archivers[format]
if fn == nil {
return nil, fmt.Errorf("%q format: %w", format, ErrUnsupportedArchiveFormat)
}
return fn(w, dir, level)
}
// NewExtractor returns a new Extractor of the specified format.
//
// The extractor will extract files to the directory provided.
func NewExtractor(format Format, r io.ReaderAt, size int64, dir string) (Extractor, error) {
fn := extractors[format]
if fn == nil {
return nil, fmt.Errorf("%q format: %w", format, ErrUnsupportedArchiveFormat)
}
return fn(r, size, dir)
}
package archive_test
import (
"errors"
"io/ioutil"
"testing"
"github.com/stretchr/testify/assert"
"gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive"
"gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive/gziplegacy"
_ "gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive/raw"
"gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive/ziplegacy"
)
func TestDefaultRegistration(t *testing.T) {
tests := map[archive.Format]struct {
hasArchiver, hasExtractor bool
}{
archive.Raw: {hasArchiver: true, hasExtractor: false},
archive.Gzip: {hasArchiver: true, hasExtractor: false},
archive.Zip: {hasArchiver: true, hasExtractor: true},
}
for tn, tc := range tests {
t.Run(string(tn), func(t *testing.T) {
_, err := archive.NewArchiver(tn, nil, "", archive.DefaultCompression)
if tc.hasArchiver {
assert.NoError(t, err)
} else {
assert.True(
t,
errors.Is(err, archive.ErrUnsupportedArchiveFormat),
"expected: %#v, got: %#v",
archive.ErrUnsupportedArchiveFormat,
err,
)
}
_, err = archive.NewExtractor(tn, nil, 0, "")
if tc.hasExtractor {
assert.NoError(t, err)
} else {
assert.True(
t,
errors.Is(err, archive.ErrUnsupportedArchiveFormat),
"expected: %#v, got: %#v",
archive.ErrUnsupportedArchiveFormat,
err,
)
}
})
}
}
func TestRegister(t *testing.T) {
format := archive.Format("new-format")
archive.Register(format, ziplegacy.NewArchiver, ziplegacy.NewExtractor)
_, err := archive.NewArchiver(format, nil, "", archive.DefaultCompression)
assert.NoError(t, err)
_, err = archive.NewExtractor(format, nil, 0, "")
assert.NoError(t, err)
}
func TestRegisterOverride(t *testing.T) {
existingGzipArchiver, err := gziplegacy.NewArchiver(ioutil.Discard, "", archive.DefaultCompression)
assert.NoError(t, err)
existingZipArchiver, err := ziplegacy.NewArchiver(ioutil.Discard, "", archive.DefaultCompression)
assert.NoError(t, err)
existingZipExtractor, err := ziplegacy.NewExtractor(nil, 0, "")
assert.NoError(t, err)
// assert existing archiver
archiver, err := archive.NewArchiver(archive.Gzip, nil, "", archive.DefaultCompression)
assert.NoError(t, err)
assert.IsType(t, existingGzipArchiver, archiver)
_, err = archive.NewExtractor(archive.Gzip, nil, 0, "")
assert.Error(t, err)
// override
archive.Register(archive.Gzip, ziplegacy.NewArchiver, ziplegacy.NewExtractor)
archiver, err = archive.NewArchiver(archive.Gzip, nil, "", archive.DefaultCompression)
assert.NoError(t, err)
assert.IsType(t, existingZipArchiver, archiver)
extractor, err := archive.NewExtractor(archive.Gzip, nil, 0, "")
assert.NoError(t, err)
assert.IsType(t, existingZipExtractor, extractor)
}
package gziplegacy
import (
"context"
"io"
"os"
"sort"
"gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive"
"gitlab.com/gitlab-org/gitlab-runner/helpers/archives"
)
func init() {
archive.Register(archive.Gzip, NewArchiver, nil)
}
// archiver is a gzip stream archiver.
type archiver struct {
w io.Writer
dir string
}
// NewArchiver returns a new Gzip Archiver.
func NewArchiver(w io.Writer, dir string, level archive.CompressionLevel) (archive.Archiver, error) {
return &archiver{w: w, dir: dir}, nil
}
// Archive archives all files as new gzip streams.
func (a *archiver) Archive(ctx context.Context, files map[string]os.FileInfo) error {
sorted := make([]string, 0, len(files))
for filename := range files {
sorted = append(sorted, filename)
}
sort.Strings(sorted)
return archives.CreateGzipArchive(a.w, sorted)
}
// Code generated by mockery v1.1.0. DO NOT EDIT.
package archive
import (
context "context"
os "os"
mock "github.com/stretchr/testify/mock"
)
// MockArchiver is an autogenerated mock type for the Archiver type
type MockArchiver struct {
mock.Mock
}
// Archive provides a mock function with given fields: ctx, files
func (_m *MockArchiver) Archive(ctx context.Context, files map[string]os.FileInfo) error {
ret := _m.Called(ctx, files)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, map[string]os.FileInfo) error); ok {
r0 = rf(ctx, files)
} else {
r0 = ret.Error(0)
}
return r0
}
// Code generated by mockery v1.1.0. DO NOT EDIT.
package archive
import (
context "context"
mock "github.com/stretchr/testify/mock"
)
// MockExtractor is an autogenerated mock type for the Extractor type
type MockExtractor struct {
mock.Mock
}
// Extract provides a mock function with given fields: ctx
func (_m *MockExtractor) Extract(ctx context.Context) error {
ret := _m.Called(ctx)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context) error); ok {
r0 = rf(ctx)
} else {
r0 = ret.Error(0)
}
return r0
}
package raw
import (
"context"
"errors"
"io"
"os"
"gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive"
)
func init() {
archive.Register(archive.Raw, NewArchiver, nil)
}
// ErrTooManyRawFiles is returned if more than one file is passed to the
// RawArchiver.
var ErrTooManyRawFiles = errors.New("only one file can be sent as raw")
// archiver is a raw archiver. It doesn't support compression nor multiple
// files.
type archiver struct {
w io.Writer
dir string
}
// NewArchiver returns a new Raw Archiver.
func NewArchiver(w io.Writer, dir string, level archive.CompressionLevel) (archive.Archiver, error) {
return &archiver{w: w, dir: dir}, nil
}
// Archive opens and copies a single file to the writer passed to
// NewRawArchiver. If more than one file is passed, ErrTooManyRawFiles is
// returned.
func (a *archiver) Archive(ctx context.Context, files map[string]os.FileInfo) error {
if len(files) > 1 {
return ErrTooManyRawFiles
}
for pathname := range files {
f, err := os.Open(pathname)
if err != nil {
return err
}
defer f.Close()
_, err = io.Copy(a.w, f)
return err
}
return nil
}
package ziplegacy
import (
"context"
"io"
"os"
"sort"
"gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive"
"gitlab.com/gitlab-org/gitlab-runner/helpers/archives"
)
func init() {
archive.Register(archive.Zip, NewArchiver, NewExtractor)
}
// archiver is a zip stream archiver.
type archiver struct {
w io.Writer
dir string
}
// NewArchiver returns a new Zip Archiver.
func NewArchiver(w io.Writer, dir string, level archive.CompressionLevel) (archive.Archiver, error) {
return &archiver{w: w, dir: dir}, nil
}
// Archive archives all files as new gzip streams.
func (a *archiver) Archive(ctx context.Context, files map[string]os.FileInfo) error {
sorted := make([]string, 0, len(files))
for filename := range files {
sorted = append(sorted, filename)
}
sort.Strings(sorted)
return archives.CreateZipArchive(a.w, sorted)
}
package ziplegacy
import (
"archive/zip"
"context"
"io"
"gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive"
"gitlab.com/gitlab-org/gitlab-runner/helpers/archives"
)
// extractor is a zip stream extractor.
type extractor struct {
r io.ReaderAt
size int64
dir string
}
// NewExtractor returns a new Zip Extractor.
func NewExtractor(r io.ReaderAt, size int64, dir string) (archive.Extractor, error) {
return &extractor{r: r, size: size, dir: dir}, nil
}
// Extract extracts files from the reader to the directory passed to
// NewZipExtractor.
func (e *extractor) Extract(ctx context.Context) error {
zr, err := zip.NewReader(e.r, e.size)
if err != nil {
return err
}
return archives.ExtractZipArchive(zr)
}
package helpers
import (
// auto-register default archivers/extractors
_ "gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive/gziplegacy"
_ "gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive/raw"
_ "gitlab.com/gitlab-org/gitlab-runner/commands/helpers/archive/ziplegacy"
)
func init() {
// Register archivers/extractors based on feature flags/environment:
// https://gitlab.com/gitlab-org/gitlab-runner/-/merge_requests/2210
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment