Hochwasser/main.go

package main

import (
	"flag"
	"fmt"
	"image"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"runtime/pprof"
	"sync"
	"time"

	"github.com/SpeckiJ/Hochwasser/pixelflut"
	"github.com/SpeckiJ/Hochwasser/render"
	"github.com/SpeckiJ/Hochwasser/rpc"
)

var (
	imgPath      = flag.String("image", "", "Filepath of an image to flut")
	ránAddr      = flag.String("rán", "", "Start RPC server to distribute jobs, listening on the given address/port")
	hevringAddr  = flag.String("hevring", "", "Connect to PRC server at given address/port")
	address      = flag.String("host", ":1234", "Target server address")
	connections  = flag.Int("connections", 4, "Number of simultaneous connections. Each connection posts a subimage")
	x            = flag.Int("x", 0, "Offset of posted image from left border")
	y            = flag.Int("y", 0, "Offset of posted image from top border")
	fetchImgPath = flag.String("fetch", "", "Enable fetching the screen area to the given local file, updating it each second")
	cpuprofile   = flag.String("cpuprofile", "", "Destination file for CPU Profile")
)

func main() {
	flag.Parse()
	rand.Seed(time.Now().UnixNano())
	task := runWithExitHandler(taskFromFlags)
	if *cpuprofile != "" {
		runWithProfiler(*cpuprofile, task)
	} else {
		task()
	}
}

func taskFromFlags(stop chan bool, wg *sync.WaitGroup) {
	rán := *ránAddr
	hev := *hevringAddr

	startServer := rán != "" || (hev == "" && *imgPath != "")
	startClient := hev != "" || (rán == "" && *imgPath != "")
	fetchImg := *fetchImgPath != ""

	if !(startServer || startClient || fetchImg) {
		fmt.Println("Error: At least one of the following flags is needed:\n	-image -rán -hevring\n")
		flag.Usage()
		os.Exit(1)
	}

	if startServer && startClient && rán == "" && hev == "" {
		rán = fmt.Sprintf(":%d", rand.Intn(30000)+1000)
		hev = rán
	}

	if startServer {
		r := rpc.SummonRán(rán, stop, wg)

		var img *image.NRGBA
		if *imgPath != "" {
			var err error
			if img, err = render.ReadImage(*imgPath); err != nil {
				log.Fatal(err)
			}
		}

		r.SetTask(pixelflut.FlutTask{
			FlutTaskOpts: pixelflut.FlutTaskOpts{
				Address:  *address,
				MaxConns: *connections,
				Offset:   image.Pt(*x, *y),
				Shuffle:  true,
			},
			Img: img,
		})
	}

	if startClient {
		rpc.ConnectHevring(hev, stop, wg)
	}

	if fetchImg {
		canvasToFile(*fetchImgPath, *address, time.Second, stop, wg)
	}
}

func canvasToFile(filepath, server string, interval time.Duration, stop chan bool, wg *sync.WaitGroup) {
	// async fetch the image
	fetchedImg := pixelflut.FetchImage(nil, server, 1, stop)

	// write it in a fixed interval
	go func() {
		wg.Add(1)
		defer wg.Done()

		for loop := true; loop; {
			select {
			case <-stop:
				loop = false
			case <-time.Tick(interval):
			}
			render.WriteImage(filepath, fetchedImg)
		}
	}()
}

// Takes a non-blocking function, and provides it an interface for graceful shutdown:
// stop chan is closed if the routine should be stopped. before quitting, wg is awaited.
func runWithExitHandler(task func(stop chan bool, wg *sync.WaitGroup)) func() {
	return func() {
		wg := sync.WaitGroup{}
		stopChan := make(chan bool)
		interruptChan := make(chan os.Signal)
		signal.Notify(interruptChan, os.Interrupt)

		task(stopChan, &wg)

		// block until we get an interrupt, or somebody says we need to quit (by closing stopChan)
		select {
		case <-interruptChan:
		case <-stopChan:
			stopChan = nil
		}

		if stopChan != nil {
			// notify all async tasks to stop on interrupt, if channel wasn't closed already
			close(stopChan)
		}

		// then wait for clean shutdown of all tasks before exiting
		wg.Wait()
	}
}

func runWithProfiler(outfile string, task func()) {
	f, err := os.Create(outfile)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	pprof.StartCPUProfile(f)
	defer pprof.StopCPUProfile()
	task()
}