github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/tools/blocksconvert/scanner/files.go (about)

     1  package scanner
     2  
     3  import (
     4  	"encoding/json"
     5  	"io"
     6  	"os"
     7  	"path/filepath"
     8  	"sync"
     9  
    10  	"github.com/golang/snappy"
    11  	"github.com/prometheus/client_golang/prometheus"
    12  	tsdb_errors "github.com/prometheus/prometheus/tsdb/errors"
    13  )
    14  
    15  type file struct {
    16  	mu   sync.Mutex
    17  	file *os.File
    18  	comp io.WriteCloser
    19  	enc  *json.Encoder
    20  }
    21  
    22  // Provides serialized access to writing entries.
    23  type openFiles struct {
    24  	mu    sync.Mutex
    25  	files map[string]*file
    26  
    27  	openFiles prometheus.Gauge
    28  }
    29  
    30  func newOpenFiles(openFilesGauge prometheus.Gauge) *openFiles {
    31  	of := &openFiles{
    32  		files:     map[string]*file{},
    33  		openFiles: openFilesGauge,
    34  	}
    35  
    36  	return of
    37  }
    38  
    39  func (of *openFiles) appendJSONEntryToFile(dir, filename string, data interface{}, headerFn func() interface{}) error {
    40  	f, err := of.getFile(dir, filename, headerFn)
    41  	if err != nil {
    42  		return err
    43  	}
    44  
    45  	// To avoid mixed output from different writes, make sure to serialize access to the file.
    46  	f.mu.Lock()
    47  	defer f.mu.Unlock()
    48  	return f.enc.Encode(data)
    49  }
    50  
    51  func (of *openFiles) getFile(dir, filename string, headerFn func() interface{}) (*file, error) {
    52  	of.mu.Lock()
    53  	defer of.mu.Unlock()
    54  
    55  	name := filepath.Join(dir, filename+".snappy")
    56  
    57  	f := of.files[name]
    58  	if f == nil {
    59  		err := os.MkdirAll(dir, os.FileMode(0700))
    60  		if err != nil {
    61  			return nil, err
    62  		}
    63  
    64  		fl, err := os.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
    65  		if err != nil {
    66  			return nil, err
    67  		}
    68  
    69  		comp := snappy.NewBufferedWriter(fl)
    70  		enc := json.NewEncoder(comp)
    71  		enc.SetEscapeHTML(false)
    72  
    73  		if headerFn != nil {
    74  			err := enc.Encode(headerFn())
    75  			if err != nil {
    76  				_ = fl.Close()
    77  			}
    78  		}
    79  
    80  		f = &file{
    81  			file: fl,
    82  			comp: comp,
    83  			enc:  enc,
    84  		}
    85  		of.files[name] = f
    86  		of.openFiles.Inc()
    87  	}
    88  
    89  	return f, nil
    90  }
    91  
    92  func (of *openFiles) closeAllFiles(footerFn func() interface{}) error {
    93  	of.mu.Lock()
    94  	defer of.mu.Unlock()
    95  
    96  	errs := tsdb_errors.NewMulti()
    97  
    98  	for fn, f := range of.files {
    99  		delete(of.files, fn)
   100  		of.openFiles.Dec()
   101  
   102  		if footerFn != nil {
   103  			errs.Add(f.enc.Encode(footerFn()))
   104  		}
   105  
   106  		errs.Add(f.comp.Close())
   107  		errs.Add(f.file.Close())
   108  	}
   109  
   110  	return errs.Err()
   111  }