143 lines
3.6 KiB
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()
|
|
}
|