github.com/kubeshop/testkube@v1.17.23/cmd/tcl/testworkflow-toolkit/artifacts/tarcached_processor.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  	"fmt"
    13  	"io"
    14  	"io/fs"
    15  	"os"
    16  	"sync"
    17  
    18  	"github.com/dustin/go-humanize"
    19  
    20  	"github.com/kubeshop/testkube/pkg/tmp"
    21  	"github.com/kubeshop/testkube/pkg/ui"
    22  )
    23  
    24  func NewTarCachedProcessor(name string, cachePath string) Processor {
    25  	if cachePath == "" {
    26  		cachePath = tmp.Name()
    27  	}
    28  	return &tarCachedProcessor{
    29  		name:      name,
    30  		cachePath: cachePath,
    31  	}
    32  }
    33  
    34  type tarCachedProcessor struct {
    35  	uploader  Uploader
    36  	name      string
    37  	cachePath string
    38  	mu        *sync.Mutex
    39  	errCh     chan error
    40  	file      *os.File
    41  	ts        *tarStream
    42  }
    43  
    44  func (d *tarCachedProcessor) Start() (err error) {
    45  	d.errCh = make(chan error)
    46  	d.mu = &sync.Mutex{}
    47  	d.file, err = os.Create(d.cachePath)
    48  
    49  	return err
    50  }
    51  
    52  func (d *tarCachedProcessor) init(uploader Uploader) {
    53  	if d.ts != nil {
    54  		return
    55  	}
    56  	d.ts = NewTarStream()
    57  	d.uploader = uploader
    58  	go func() {
    59  		_, err := io.Copy(d.file, d.ts)
    60  		d.errCh <- err
    61  	}()
    62  }
    63  
    64  func (d *tarCachedProcessor) clean() {
    65  	_ = os.Remove(d.cachePath)
    66  }
    67  
    68  func (d *tarCachedProcessor) upload(path string, file fs.File, stat fs.FileInfo) error {
    69  	defer file.Close()
    70  	return d.ts.Add(path, file, stat)
    71  }
    72  
    73  func (d *tarCachedProcessor) Add(uploader Uploader, path string, file fs.File, stat fs.FileInfo) error {
    74  	d.mu.Lock()
    75  	d.init(uploader)
    76  	defer d.mu.Unlock()
    77  	return d.upload(path, file, stat)
    78  }
    79  
    80  func (d *tarCachedProcessor) End() (err error) {
    81  	defer d.clean()
    82  
    83  	if d.ts != nil {
    84  		<-d.ts.Done()
    85  	}
    86  	err = d.ts.Close()
    87  	if err != nil {
    88  		return fmt.Errorf("problem closing writer: %w", err)
    89  	}
    90  	err = <-d.errCh
    91  	if err != nil {
    92  		return fmt.Errorf("problem writing to disk cache: %w", err)
    93  	}
    94  
    95  	if d.uploader == nil {
    96  		return nil
    97  	}
    98  
    99  	file, err := os.Open(d.cachePath)
   100  	if err != nil {
   101  		return fmt.Errorf("problem reading disk cache: %w", err)
   102  	}
   103  
   104  	stat, err := file.Stat()
   105  	if err != nil {
   106  		return fmt.Errorf("problem reading disk cache: stat: %w", err)
   107  	}
   108  
   109  	fmt.Printf("Archived everything in %s archive (%s).\n", ui.LightCyan(d.name), ui.LightCyan(humanize.Bytes(uint64(stat.Size()))))
   110  	return d.uploader.Add(d.name, file, stat.Size())
   111  }