split.go

package main

import (
	"bufio"
	"bytes"
	"flag"
	"io"
	"log"
	"os"
	"strings"
)

func main() {
	chunkLength := flag.Uint64("cl", 1, "chunk length")
	fileSize := flag.Uint64("fs", 10485760, "file size")
	recordSeparatorStr := flag.String("rs", "\n", "record separator (single byte)")
	flag.Parse()
	if len(*recordSeparatorStr) != 1 {
		log.Fatal("record separator contains multiple bytes")
	}
	recordSeparator := []byte(*recordSeparatorStr)[0]
	prefix := flag.Arg(0)
	if prefix == "" {
		log.Fatal("no prefix supplied")
	}

	r := bufio.NewReader(os.Stdin)
	f, err := file(prefix, 0)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	w := bufio.NewWriter(f)

	var eof bool
	var index int
	var bytesWritten uint64
	for {
		var chunk bytes.Buffer
		for i := uint64(0); i < *chunkLength; i++ {
			record, err := r.ReadBytes(recordSeparator)
			if err != nil {
				if err != io.EOF {
					log.Fatal(err)
				}
				eof = true
			} else {
				chunk.Write(record)
			}
		}

		chunkLen := uint64(chunk.Len())
		if chunkLen+bytesWritten > *fileSize && chunkLen <= *fileSize {
			index = index + 1
			if err := w.Flush(); err != nil {
				log.Fatal(err)
			}
			if err := f.Close(); err != nil {
				log.Fatal(err)
			}
			f, err = file(prefix, index)
			if err != nil {
				log.Fatal(err)
			}
			w = bufio.NewWriter(f)
			bytesWritten = 0
		}

		chunk.WriteTo(w)
		bytesWritten = bytesWritten + chunkLen

		if eof {
			break
		}
	}

	w.Flush()
}

func file(prefix string, index int) (*os.File, error) {
	z := index / 26
	return os.Create(prefix + strings.Repeat("z", z) + string(97+index-z*26))
}