Hochwasser/main.go

143 lines
3.6 KiB
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(img, image.Pt(*x, *y), *address, *connections)
}
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()
}