github.com/kubeshop/testkube@v1.17.23/cmd/tcl/testworkflow-toolkit/artifacts/direct_uploader.go (about)

     1  // Copyright 2024 Testkube.
     2  //
     3  // Licensed as a Testkube Pro file under the Testkube Community
     4  // License (the "License"); you may not use this file except in compliance with
     5  // the License. You may obtain a copy of the License at
     6  //
     7  //	https://github.com/kubeshop/testkube/blob/main/licenses/TCL.txt
     8  
     9  package artifacts
    10  
    11  import (
    12  	"context"
    13  	"fmt"
    14  	"io"
    15  	"sync"
    16  	"sync/atomic"
    17  
    18  	minio2 "github.com/minio/minio-go/v7"
    19  
    20  	"github.com/kubeshop/testkube/cmd/tcl/testworkflow-toolkit/env"
    21  	"github.com/kubeshop/testkube/pkg/storage/minio"
    22  	"github.com/kubeshop/testkube/pkg/ui"
    23  )
    24  
    25  type PutObjectOptionsEnhancer = func(options *minio2.PutObjectOptions, path string, size int64)
    26  
    27  func NewDirectUploader(opts ...DirectUploaderOpt) Uploader {
    28  	uploader := &directUploader{
    29  		parallelism: 1,
    30  		options:     make([]PutObjectOptionsEnhancer, 0),
    31  	}
    32  	for _, opt := range opts {
    33  		opt(uploader)
    34  	}
    35  	return uploader
    36  }
    37  
    38  type directUploader struct {
    39  	client      *minio.Client
    40  	wg          sync.WaitGroup
    41  	sema        chan struct{}
    42  	parallelism int
    43  	error       atomic.Bool
    44  	options     []PutObjectOptionsEnhancer
    45  }
    46  
    47  func (d *directUploader) Start() (err error) {
    48  	d.client, err = env.ObjectStorageClient()
    49  	d.sema = make(chan struct{}, d.parallelism)
    50  	return err
    51  }
    52  
    53  func (d *directUploader) buildOptions(path string, size int64) (options minio2.PutObjectOptions) {
    54  	for _, enhance := range d.options {
    55  		enhance(&options, path, size)
    56  	}
    57  	if options.ContentType == "" {
    58  		options.ContentType = "application/octet-stream"
    59  	}
    60  	return options
    61  }
    62  
    63  func (d *directUploader) upload(path string, file io.ReadCloser, size int64) {
    64  	ns := env.ExecutionId()
    65  	opts := d.buildOptions(path, size)
    66  	err := d.client.SaveFileDirect(context.Background(), ns, path, file, size, opts)
    67  
    68  	if err != nil {
    69  		d.error.Store(true)
    70  		ui.Errf("%s: failed: %s", path, err.Error())
    71  		return
    72  	}
    73  }
    74  
    75  func (d *directUploader) Add(path string, file io.ReadCloser, size int64) error {
    76  	d.wg.Add(1)
    77  	d.sema <- struct{}{}
    78  	go func() {
    79  		d.upload(path, file, size)
    80  		_ = file.Close()
    81  		d.wg.Done()
    82  		<-d.sema
    83  	}()
    84  	return nil
    85  }
    86  
    87  func (d *directUploader) End() error {
    88  	d.wg.Wait()
    89  	if d.error.Load() {
    90  		return fmt.Errorf("upload failed")
    91  	}
    92  	return nil
    93  }
    94  
    95  type DirectUploaderOpt = func(uploader *directUploader)
    96  
    97  func WithParallelism(parallelism int) DirectUploaderOpt {
    98  	return func(uploader *directUploader) {
    99  		if parallelism < 1 {
   100  			parallelism = 1
   101  		}
   102  		uploader.parallelism = parallelism
   103  	}
   104  }
   105  
   106  func WithMinioOptionsEnhancer(fn PutObjectOptionsEnhancer) DirectUploaderOpt {
   107  	return func(uploader *directUploader) {
   108  		uploader.options = append(uploader.options, fn)
   109  	}
   110  }