erigon-pulse/common/etl/dataprovider.go

package etl

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"runtime"

	"github.com/ledgerwatch/erigon/common"
	"github.com/ledgerwatch/erigon/log"
)

type dataProvider interface {
	Next(decoder Decoder) ([]byte, []byte, error)
	Dispose() (uint64, error)
}

type fileDataProvider struct {
	file   *os.File
	reader io.Reader
}

type Encoder interface {
	Encode(toWrite interface{}) error
	Reset(writer io.Writer)
}

func FlushToDisk(encoder Encoder, currentKey []byte, b Buffer, tmpdir string) (dataProvider, error) {
	if b.Len() == 0 {
		return nil, nil
	}
	// if we are going to create files in the system temp dir, we don't need any
	// subfolders.
	if tmpdir != "" {
		if err := os.MkdirAll(tmpdir, 0755); err != nil {
			return nil, err
		}
	}

	bufferFile, err := ioutil.TempFile(tmpdir, "tg-sync-sortable-buf")
	if err != nil {
		return nil, err
	}
	defer bufferFile.Sync() //nolint:errcheck

	w := bufio.NewWriterSize(bufferFile, BufIOSize)
	defer w.Flush() //nolint:errcheck

	defer func() {
		b.Reset() // run it after buf.flush and file.sync
		var m runtime.MemStats
		runtime.ReadMemStats(&m)
		log.Info(
			"Flushed buffer file",
			"name", bufferFile.Name(),
			"alloc", common.StorageSize(m.Alloc), "sys", common.StorageSize(m.Sys))
	}()

	encoder.Reset(w)
	for _, entry := range b.GetEntries() {
		err = writeToDisk(encoder, entry.key, entry.value)
		if err != nil {
			return nil, fmt.Errorf("error writing entries to disk: %v", err)
		}
	}

	return &fileDataProvider{bufferFile, nil}, nil
}

func (p *fileDataProvider) Next(decoder Decoder) ([]byte, []byte, error) {
	if p.reader == nil {
		_, err := p.file.Seek(0, 0)
		if err != nil {
			return nil, nil, err
		}
		p.reader = bufio.NewReaderSize(p.file, BufIOSize)
	}
	decoder.Reset(p.reader)
	return readElementFromDisk(decoder)
}

func (p *fileDataProvider) Dispose() (uint64, error) {
	info, errStat := os.Stat(p.file.Name())
	errClose := p.file.Close()
	errRemove := os.Remove(p.file.Name())
	if errClose != nil {
		return 0, errClose
	}
	if errRemove != nil {
		return 0, errRemove
	}
	if errStat != nil {
		return 0, errStat
	}
	return uint64(info.Size()), nil
}

func (p *fileDataProvider) String() string {
	return fmt.Sprintf("%T(file: %s)", p, p.file.Name())
}

func writeToDisk(encoder Encoder, key []byte, value []byte) error {
	toWrite := [][]byte{key, value}
	return encoder.Encode(toWrite)
}

func readElementFromDisk(decoder Decoder) ([]byte, []byte, error) {
	result := make([][]byte, 2)
	err := decoder.Decode(&result)
	return result[0], result[1], err
}

type memoryDataProvider struct {
	buffer       Buffer
	currentIndex int
}

func KeepInRAM(buffer Buffer) dataProvider {
	return &memoryDataProvider{buffer, 0}
}

func (p *memoryDataProvider) Next(decoder Decoder) ([]byte, []byte, error) {
	if p.currentIndex >= p.buffer.Len() {
		return nil, nil, io.EOF
	}
	entry := p.buffer.Get(p.currentIndex)
	p.currentIndex++
	return entry.key, entry.value, nil
}

func (p *memoryDataProvider) Dispose() (uint64, error) {
	return 0 /* doesn't take space on disk */, nil
}

func (p *memoryDataProvider) String() string {
	return fmt.Sprintf("%T(buffer.Len: %d)", p, p.buffer.Len())
}