metrics/lib/filestream/parallel.go

package filestream

import (
	"io"
	"sync"

	"github.com/VictoriaMetrics/VictoriaMetrics/lib/fs/fsutil"
	"github.com/VictoriaMetrics/VictoriaMetrics/lib/logger"
)

// ParallelFileCreator is used for parallel creating of files for the given dstPath.
//
// ParallelFileCreator is needed for speeding up creating many files on high-latency
// storage systems such as NFS or Ceph.
type ParallelFileCreator struct {
	tasks []parallelFileCreatorTask
}

type parallelFileCreatorTask struct {
	dstPath string
	wc      *WriteCloser
	nocache bool
}

// Add registers a task for creating the file at dstPath and assigning it to *wc.
//
// Tasks are executed in parallel on Run() call.
func (pfc *ParallelFileCreator) Add(dstPath string, wc *WriteCloser, nocache bool) {
	pfc.tasks = append(pfc.tasks, parallelFileCreatorTask{
		dstPath: dstPath,
		wc:      wc,
		nocache: nocache,
	})
}

// Run runs all the registered tasks for creating files in parallel.
func (pfc *ParallelFileCreator) Run() {
	var wg sync.WaitGroup
	concurrencyCh := fsutil.GetConcurrencyCh()
	for _, task := range pfc.tasks {
		concurrencyCh <- struct{}{}
		wg.Go(func() {
			*task.wc = MustCreate(task.dstPath, task.nocache)
			<-concurrencyCh
		})
	}
	wg.Wait()
}

// ParallelFileOpener is used for parallel opening of files at the given dstPath.
//
// ParallelFileOpener is needed for speeding up opening many files on high-latency
// storage systems such as NFS or Ceph.
type ParallelFileOpener struct {
	tasks []parallelFileOpenerTask
}

type parallelFileOpenerTask struct {
	path    string
	rc      *ReadCloser
	nocache bool
}

// Add registers a task for opening the file ath the given path and assigning it to *rc.
//
// Tasks are executed in parallel on Run() call.
func (pfo *ParallelFileOpener) Add(path string, rc *ReadCloser, nocache bool) {
	pfo.tasks = append(pfo.tasks, parallelFileOpenerTask{
		path:    path,
		rc:      rc,
		nocache: nocache,
	})
}

// Run runs all the registered tasks for opening files in parallel.
func (pfo *ParallelFileOpener) Run() {
	var wg sync.WaitGroup
	concurrencyCh := fsutil.GetConcurrencyCh()
	for _, task := range pfo.tasks {
		concurrencyCh <- struct{}{}
		wg.Go(func() {
			*task.rc = MustOpen(task.path, task.nocache)
			<-concurrencyCh
		})
	}
	wg.Wait()
}

// ParallelStreamWriter is used for parallel writing of data from io.WriterTo to the given dstPath files.
//
// ParallelStreamWriter is needed for speeding up writing data to many files on high-latency
// storage systems such as NFS or Ceph.
type ParallelStreamWriter struct {
	tasks []parallelStreamWriterTask
}

type parallelStreamWriterTask struct {
	dstPath string
	src     io.WriterTo
}

// Add adds a task to execute in parallel - to write the data from src to the dstPath.
//
// Tasks are executed in parallel on Run() call.
func (psw *ParallelStreamWriter) Add(dstPath string, src io.WriterTo) {
	psw.tasks = append(psw.tasks, parallelStreamWriterTask{
		dstPath: dstPath,
		src:     src,
	})
}

// Run executes all the tasks added via Add() call in parallel.
func (psw *ParallelStreamWriter) Run() {
	var wg sync.WaitGroup
	concurrencyCh := fsutil.GetConcurrencyCh()
	for _, task := range psw.tasks {
		concurrencyCh <- struct{}{}

		wg.Go(func() {
			f := MustCreate(task.dstPath, false)
			if _, err := task.src.WriteTo(f); err != nil {
				f.MustClose()
				// Do not call MustRemovePath(path), so the user could inspect
				// the file contents during investigation of the issue.
				logger.Panicf("FATAL: cannot write data to %q: %s", task.dstPath, err)
			}
			f.MustClose()

			<-concurrencyCh
		})
	}
	wg.Wait()
}