Untitled

package main

import (
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Range struct {
	start int
	end   int
}

type Params struct {
	file    string
	timeout int
	ranges  []Range
}

type SyncWriter struct {
	mutex  sync.Mutex
	Writer io.Writer
}

func (writer *SyncWriter) Write(data []byte) (n int, err error) {
	writer.mutex.Lock()
	defer writer.mutex.Unlock()
	return writer.Writer.Write(data)
}

func getArgs(args []string, params *Params) error {
	var err error
	for i := range args {
		if args[i][:2] == "--" && i != len(args)-1 {
			arg := args[i][2:]
			body := args[i+1]
			switch arg {
			case "file":
				params.file = body
			case "timeout":
				params.timeout, err = strconv.Atoi(body)
				if err != nil {
					return err
				}
			case "range":
				rng := strings.Split(body, ":")
				start, err := strconv.Atoi(rng[0])
				if err != nil {
					return err
				}
				end, err := strconv.Atoi(rng[1])
				if err != nil {
					return err
				}
				params.ranges = append(params.ranges, Range{
					start: start,
					end:   end,
				})
			}
		}
	}
	return err
}

func primeNums(ch chan []int, rng Range) {
	var nums []int
	for i := rng.start; i <= rng.end; i++ {
		var fl bool
		for j := 2; j < i; j++ {
			if i%j == 0 {
				fl = true
				break
			}
		}
		if !fl {
			nums = append(nums, i)
		}
	}
	ch <- nums
}

func saveNums(wg *sync.WaitGroup, ch chan []int, file *os.File) {
	nums := <-ch
	for _, num := range nums {
		fmt.Fprintf(file, "%v\n", num)
	}
	wg.Done()
}

func calc(ctx context.Context, file *os.File, ranges []Range) {

	wg := sync.WaitGroup{}
	exit := make(chan struct{})

	wg.Add(len(ranges))
	go func() {
		for _, rng := range ranges {
			ch := make(chan []int)
			go primeNums(ch, rng)
			go saveNums(&wg, ch, file)
		}
		wg.Wait()
		close(exit)
	}()

	for {
		select {
		case <-ctx.Done():
			log.Printf("end: %v", ctx.Err())
			return
		case <-exit:
			log.Println("END with exit")
			return
		default:
			log.Println("running")
			time.Sleep(time.Second)
		}
	}
}

func main() {
	var params *Params = &Params{}

	err := getArgs(os.Args[1:], params)
	if err != nil {
		panic(err)
	}

	file, err := os.Create(params.file)
	if err != nil {
		panic(err)
	}
	defer file.Close()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second*time.Duration(params.timeout))
	defer cancel()

	calc(ctx, file, params.ranges)
}