scatter.go

package gopar3

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/klauspost/reedsolomon"
	"golang.org/x/sync/errgroup"
)

// Scatter acts the same as [Inflate], but outputs each shard
// as its own file. Recovery is the same, as long as the quorum
// of files is available.
func Scatter(
	ctx context.Context,
	destination string,
	source string,
	withOptions ...Option,
) (err error) {
	options, err := newOptions(withOptions...)
	if err != nil {
		return err
	}

	f, err := os.Stat(destination)
	if err != nil {
		return err
	}
	if !f.IsDir() {
		return errors.New("destination must be a directory")
	}
	f, err = os.Stat(source)
	if err != nil {
		return err
	}
	if f.IsDir() {
		return errors.New("cannot inflate a directory")
	}
	r, err := os.Open(source)
	if err != nil {
		return err
	}
	defer func() {
		err = errors.Join(err, r.Close())
	}()

	var (
		l = &BatchLoader{
			Quorum:    int(options.QuorumShards),
			Shards:    int(options.QuorumShards + options.RedundantShards),
			ShardSize: options.ShardSize,
		}
	)
	rs, err := reedsolomon.New(
		l.Quorum, l.Shards-l.Quorum,
		reedsolomon.WithAutoGoroutines(options.ShardSize),
	)
	if err != nil {
		return err
	}

	tag, err := NewTag(ctx, r, options.QuorumShards)
	if err != nil {
		return err
	}
	if _, err = r.Seek(0, io.SeekStart); err != nil {
		return err
	}

	wg, ctx := errgroup.WithContext(ctx)
	batches := make(chan [][]byte, 4)
	wg.Go(func() (err error) {
		defer close(batches)
		var batch [][]byte
		for {
			batch, _, err = l.Load(r)
			if err != nil {
				// panic(fmt.Sprintf("%+v", batch))
				if err == io.EOF {
					return nil
				}
				return err
			}
			select {
			case <-ctx.Done():
				return ctx.Err()
			case batches <- batch:
			}
		}
	})

	batchesWithParity := make(chan [][]byte, 4)
	wg.Go(func() (err error) {
		defer close(batchesWithParity)
		for batch := range batches {
			if err = rs.Reconstruct(batch); err != nil {
				return err
			}
			select {
			case <-ctx.Done():
				return ctx.Err()
			case batchesWithParity <- batch:
			}
		}
		return nil
	})

	splitShards := make([]chan []byte, l.Shards)
	for i := range splitShards {
		splitShards[i] = make(chan []byte, 4)
	}
	wg.Go(func() (err error) {
		for batch := range batchesWithParity {
			for i, shard := range batch {
				select {
				case <-ctx.Done():
					return ctx.Err()
				case splitShards[i] <- shard:
				}
			}
		}

		for i := range splitShards {
			close(splitShards[i])
		}
		return nil
	})

	ext := filepath.Ext(source)
	base := strings.TrimSuffix(filepath.Base(source), ext)
	lastBatch := tag.Batches(int64(options.ShardSize)) - 1
	for i := range splitShards {
		i := i
		wg.Go(func() (err error) {
			splitTag := tag
			splitTag.ShardOrder = uint8(i)
			for shard := range splitShards[i] {
				if err = writeOneShard(
					filepath.Join(
						destination,
						fmt.Sprintf(
							`%s%x%s.%dof%dof%d.gopar3`,
							base,
							tag.SourceCRC,
							ext,
							i, splitTag.ShardBatch, lastBatch,
						),
					),
					splitTag.Bytes(),
					shard,
					options.Telomeres,
				); err != nil {
					return err
				}
				splitTag.ShardBatch++
			}
			return nil
		})
	}
	return wg.Wait()
}

func writeOneShard(destination string, tag, shard []byte, telomereCount int) (err error) {
	w, closer, err := NewShardWriter(
		destination,
		telomereCount,
		nil,
	)
	if err != nil {
		return err
	}
	defer func() {
		err = errors.Join(err, closer())
	}()
	return w.Write(tag, shard)
}