cyclescaler/master/main.go

package main

import (
	"bytes"
	"crypto/md5"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"html/template"
	"image/png"
	"io"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"time"
	"unicode"

	"deadbeef.codes/steven/cyclescaler"
	"github.com/disintegration/imaging"
)

// Wip (work in progress) is instanciated for each frame / task that a slave is currently working on.  They are destroyed when a slave is finished or if we've lost contact with the slave.
type Wip struct {
	FrameNumber        int
	HeartbeatTimestamp time.Time
	Output             string
	Timeout            bool // Haven't received a hearbeat and slated for cleanup
	CompletePercentage int
	SlaveName          string // name of the slave processing this frame
	BadHealth          bool   // If we're missing heartbeats
}

// BatchJob is instanciated for each .blend file uploaded and contains data related to the total rendering that needs to happen
// The BatchJob contains the state of all the frames in the batch job.  A frame can be in 3 different states.  It can
// be sitting in RemainingFrames, it can be in a Wip, or it can be in CompletedFrames.  When it moves to one state,
// it is removed from the previous one.  If a job times out, the watchdog handler will throw it back into the remaining
// frames state
// We use 3 slices for tracking state instead of keeping the state inside of a "frame" object because it scales better when
// searching for frames in a specific state
type BatchJob struct {
	Token         string // the token for the blend file
	BlendFileName string // original blend file name

	//Job state
	Started            bool // whether the job is paused or not yet started
	Completed          bool // whether the job is completed or not
	CompletePercentage int  //

	//Frame state
	RemainingFrames []int // these are "undispatched" frames.  If it's in progress, it's not in here, doesn't mean it's rendered.
	Wips            []Wip // work in progress, frames which have been scheduled to be rendered by a slave.  Contains heartbeat timestamp
	CompletedFrames []int // frames which have been completed

	NumberOfFrames int // Always equal to len(RemainingFrames) + len(Wips) + len (CompletedFrames)
	StartTime      string
}

// HomePage is data that gets served to templates/home
type HomePage struct {
	NotNoJobs bool // This is purposely a dumb name.  If there are any previously or ongoing jobs... AKA this is only false if this master has just come online and nothing has ever been uploaded to it.  If it's true, the home page will display the segment listing ongoing jobs.
	Jobs      *map[string]*BatchJob
	Token     string // This is the token that's used if a new batch job were to be uploaded
}

// JobStatusPageImages is instantiated for each image on the status page which has a Thumbnail
// as well as a RenderedFramePath.  This struct maps the two together.
type JobStatusPageImages struct {
	ThumbnailPath     string
	RenderedFramePath string
}

// JobStatusPageData that is served to templates/jobStatus.html
type JobStatusPage struct {
	Images         []JobStatusPageImages
	UseTinyImages  bool //If we have more than 100 images
	Job            *BatchJob
	CompleteFrames string
	LatestImage    string
	InProgress     bool
}

// Data that is served to templates/invalidMime.html
type InvalidMimePage struct {
	UploadedFileMime string
}

var (
	batchJobs map[string]*BatchJob // Uses session token as key

)

// http server setup and routing
func main() {

	batchJobs = make(map[string]*BatchJob)
	go watchdogTimer()

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt)
	go func() {
		http.Handle("/public/", http.StripPrefix("/public/", http.FileServer(http.Dir("public"))))
		http.HandleFunc("/", homePageHandler)
		http.HandleFunc("/settings", settingsPageHandler)
		http.HandleFunc("/jobStatus", jobStatusHandler)
		http.HandleFunc("/getwork", getWorkHandler)
		http.HandleFunc("/postwork", postWorkHandler)
		http.HandleFunc("/heartbeat", heartbeatHandler)
		http.HandleFunc("/404", notFoundHandler)
		log.Fatal(http.ListenAndServe(":8097", nil))
	}()
	log.Println("started cyclescaler on 8097")

	<-stop
	fmt.Println("Shutting server down...")
}

// Collect the .blend file from the user
// Display list of running jobs
func homePageHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" {
		timeNow := time.Now().Unix()
		h := md5.New()
		io.WriteString(h, strconv.FormatInt(timeNow, 10))

		homePage := HomePage{}

		fmt.Println(len(batchJobs))
		if len(batchJobs) > 0 {
			homePage.NotNoJobs = true
			homePage.Jobs = &batchJobs
		}

		homePage.Token = fmt.Sprintf("%x", h.Sum(nil))

		t, err := template.ParseFiles("templates/home.html")
		if err != nil {
			log.Fatalf("could not parse home.html template file: %v", err)
		}
		err = t.Execute(w, homePage)
		if err != nil {
			log.Fatalf("could not execute home.html template file: %v", err)
		}

	} else {
		err := r.ParseMultipartForm(32 << 22)
		if err != nil {
			log.Printf("failed to parse form during file upload: %v", err)
			return
		}
		token, ok := r.Form["token"]
		if !ok {
			log.Printf("couldn't GET token")
			return
		}

		file, handler, err := r.FormFile("uploadfile")
		if err != nil {
			fmt.Println(err)
			return
		}

		defer file.Close()
		if handler.Header.Get("Content-Type") != "application/x-blender" && handler.Header.Get("Content-Type") != "application/octet-stream" {
			t, err := template.ParseFiles("templates/invalidMime.html")
			if err != nil {
				log.Fatalf("could not parse invalidMime.html template file: %v", err)
			}

			invalidMimePage := &InvalidMimePage{}
			invalidMimePage.UploadedFileMime = handler.Header.Get("Content-Type")

			t.Execute(w, invalidMimePage)
			return
		}
		http.Redirect(w, r, fmt.Sprintf("/settings?token=%s", r.Form["token"][0]), http.StatusSeeOther)

		// Store the uploaded file in the public directory and run it through a SHA256 hash function
		f, err := os.OpenFile(fmt.Sprintf("public/blendfiles/%s.blend", token[0]), os.O_WRONLY|os.O_CREATE, 0666)
		if err != nil {
			log.Printf("could not create /public/blendfiles/%s.blend: %v", token[0], err)
			return
		}
		defer f.Close()
		hash := sha256.New()
		multiWriter := io.MultiWriter(f, hash)
		if _, err := io.Copy(multiWriter, file); err != nil {
			log.Printf("could not copy file to multiWriter (file output and hash engine): %v", err)
			return
		}

		// Store the output of the hash function to a file in the public directory
		hf, err := os.OpenFile(fmt.Sprintf("public/blendfiles/%s.sha256", token[0]), os.O_WRONLY|os.O_CREATE, 0666)
		if err != nil {
			log.Printf("coult not create /public/blendfiles/%s.sha256: %v", token[0], err)
			return
		}
		hr := bytes.NewReader(hash.Sum(nil))
		if _, err := io.Copy(hf, hr); err != nil {
			log.Printf("could not copy hash of uploaded blend file to /public/blendfiles/%s.sha256: %v", token[0], err)
			return
		}

		t := time.Now()

		// Create the job and set known parameters
		batchJob := BatchJob{
			BlendFileName: handler.Filename,
			Token:         r.Form["token"][0],
			StartTime:     t.Format(time.RFC3339)}
		batchJobs[r.Form["token"][0]] = &batchJob
	}
}

// Collect settings from the user for the batch job
func settingsPageHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" {
		err := r.ParseForm()
		if err != nil {
			log.Printf("failed to parse form in settingsHandler GET: %v", err)
			return
		}

		_, ok := r.Form["token"]
		if !ok {
			log.Println("couldn't get token in settingsHandler GET")
			return
		}

		t, err := template.ParseFiles("templates/settings.html")
		if err != nil {
			log.Fatalf("could not parse settings.html template file: %v", err)
		}
		t.Execute(w, batchJobs[r.Form["token"][0]])

	} else {
		// Begin render and redirect person to status page for that render
		err := r.ParseMultipartForm(32 << 20)
		if err != nil {
			log.Printf("failed to parse form in settingsHandler POST: %v", err)
			return
		}
		token, ok := r.Form["token"]
		if !ok {
			log.Println("couldn't get token in settingsHandler POST")
			return
		}

		var renderFrames string
		_, ok = r.Form["renderframes"]
		if ok {
			renderFrames = stripSpaces(r.Form["renderframes"][0]) + "," //This is a hack to trigger the , case below to add the last set of numbers
		}

		log.Printf("renderFrames = %s", renderFrames)

		//This is where we have to parse the input settings into job.RemainingFrames
		var remainingFrames []int
		var firstNumber int //Used for ranges of frames, such as 1-50 in the case matching below
		firstNumber = -1
		var currentNumber string

		for _, char := range renderFrames {
			switch char {
			case ',':
				frame, err := strconv.Atoi(currentNumber)
				log.Printf("frame = %d", frame)
				if err != nil {
					log.Printf("error converting currentNumber '%s' to integer: %v", currentNumber, err)
				}
				if firstNumber == -1 {
					remainingFrames = append(remainingFrames, frame)
				} else {
					for j := firstNumber; j < frame+1; j++ {
						remainingFrames = append(remainingFrames, j)
					}
				}
				firstNumber = -1
				currentNumber = ""

			case '-':
				firstNumber, err = strconv.Atoi(currentNumber)
				if err != nil {
					log.Printf("error converting currentNumber '%s' to integer: %v", currentNumber, err)
				}
				currentNumber = ""

			default:
				currentNumber += string(char)
			}
		}

		batchJobs[token[0]].RemainingFrames = remainingFrames
		batchJobs[token[0]].NumberOfFrames = len(remainingFrames)
		batchJobs[token[0]].Started = true

		renderedFramesDir := fmt.Sprintf("public/renderedframes/%s", token[0])
		renderedThumbnailsDir := fmt.Sprintf("public/renderedframes-thumbnail/%s", token[0])
		if _, err := os.Stat(renderedFramesDir); os.IsNotExist(err) {
			os.Mkdir(renderedFramesDir, os.ModePerm)
		}
		if _, err := os.Stat(renderedThumbnailsDir); os.IsNotExist(err) {
			os.Mkdir(renderedThumbnailsDir, os.ModePerm)
		}

		http.Redirect(w, r, fmt.Sprintf("/jobStatus?token=%s", r.Form["token"][0]), http.StatusSeeOther)
	}
}

// The status page to review status for the render job, see frame thumbnails and renders in progress
func jobStatusHandler(w http.ResponseWriter, r *http.Request) {
	err := r.ParseForm()
	if err != nil {
		log.Printf("failed to parse form in statusHandler: %v", err)
		return
	}

	_, ok := r.Form["token"]
	if !ok {
		log.Println("couldn't get token in statusHandler")
		return
	}
	token := r.Form["token"][0]

	_, ok = batchJobs[token]
	if !ok {
		log.Println("request to see status page with token that doesn't exist")
		http.Redirect(w, r, "/404", http.StatusSeeOther)
		return
	}

	t, err := template.ParseFiles("templates/jobStatus.html")
	if err != nil {
		log.Fatalf("could not parse jobStatus.html template file: %v", err)
	}

	jobStatusPage := JobStatusPage{}
	jobStatusPage.Job = batchJobs[token]

	f, err := os.Open(fmt.Sprintf("public/renderedframes-thumbnail/%s", token))
	if err != nil {
		log.Fatal(err)
	}
	files, err := f.Readdir(-1)
	f.Close()
	if err != nil {
		log.Fatal(err)
	}

	for _, file := range files {
		tnp := fmt.Sprintf("public/renderedframes-thumbnail/%s/%s", token, file.Name())
		rfp := fmt.Sprintf("public/renderedframes/%s/%s", token, file.Name())

		if strings.HasSuffix(tnp, ".png") {
			spi := JobStatusPageImages{ThumbnailPath: tnp, RenderedFramePath: rfp}
			jobStatusPage.Images = append(jobStatusPage.Images, spi)
		}
	}
	sort.Slice(jobStatusPage.Images, func(i, j int) bool {
		return jobStatusPage.Images[i].ThumbnailPath < jobStatusPage.Images[j].ThumbnailPath
	})

	if len(jobStatusPage.Images) > 0 {
		jobStatusPage.LatestImage = jobStatusPage.Images[len(jobStatusPage.Images)-1].RenderedFramePath
	} else {
		jobStatusPage.LatestImage = "public/img/blankthumbnail.png"
	}

	if len(jobStatusPage.Images) > 500 {
		jobStatusPage.UseTinyImages = true
	}

	if len(jobStatusPage.Job.Wips) > 0 {
		jobStatusPage.InProgress = true
	} else {
		jobStatusPage.InProgress = false
	}

	t.Execute(w, jobStatusPage)
}

// A slave is requesting work - find them something to do
func getWorkHandler(w http.ResponseWriter, r *http.Request) {
	err := r.ParseForm()
	if err != nil {
		log.Printf("failed to parse form in getWorkHandler: %v", err)
		return
	}

	_, ok := r.Form["slavename"]
	if !ok {
		log.Println("couldn't get slavename getWorkHandler")
		return
	}

	slaveName := r.Form["slavename"][0]

	for _, job := range batchJobs {
		if job.Completed || !job.Started { //Jobs in these states are skipped
			continue
		}
		if len(job.RemainingFrames) > 0 {
			slaveOrder := &cyclescaler.SlaveOrder{}
			slaveOrder.Token = job.Token

			slaveOrder.FrameNumber = job.RemainingFrames[0]                                   // We grab the next frame to render
			job.RemainingFrames = append(job.RemainingFrames[:0], job.RemainingFrames[1:]...) // Remove it from the remaining frames

			wip := Wip{} //and then store it all in a WIP
			wip.FrameNumber = slaveOrder.FrameNumber
			wip.HeartbeatTimestamp = time.Now()
			wip.Output = "Job dispatched to slave..."
			wip.SlaveName = slaveName
			job.Wips = append(job.Wips, wip)

			if slaveOrder.Token != "" {
				json, err := json.Marshal(slaveOrder)
				if err != nil {
					log.Printf("failed to marshal job '%s' to json: %v", job.BlendFileName, err)
					return
				}
				_, err = w.Write(json)
				if err != nil {
					log.Printf("failed to write json to http response writer for job '%s': %v", job.BlendFileName, err)
					return
				}
				return
			}
		}
		if len(job.Wips) == 0 { //The job is now complete, but not yet marked as complete
			job.Completed = true
		}
	}
	w.Write([]byte("none")) // No work available, terminate the instance
}

// Accept a rendered png frame from a slave and save it to the public directory
func postWorkHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("received post request to upload work")
	err := r.ParseMultipartForm(32 << 20)
	if err != nil {
		log.Printf("failed to parse multipart form in postWorkHandler")
		return
	}
	token, ok := r.Form["token"]
	if !ok {
		log.Printf("couldn't get token for postWorkHandler")
		return
	}

	_, ok = r.Form["framenumber"]
	if !ok {
		log.Printf("couldn't get framenumber for postWorkHandler")
		return
	}

	frameNumber, err := strconv.Atoi(r.Form["framenumber"][0])
	if err != nil {
		log.Printf("failed to convert frame number '%s' to integer: %v", r.Form["framenumber"][0], err)
		return
	}

	fmt.Printf("postWork framenumber '%d', token '%s'", frameNumber, token[0])

	file, handler, err := r.FormFile("uploadfile")
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println("1")
	defer file.Close()

	f, err := os.OpenFile(fmt.Sprintf("public/renderedframes/%s/%s", token[0], handler.Filename), os.O_WRONLY|os.O_CREATE, 0666)
	if err != nil {
		log.Printf("failed to create uploaded file locally: %v", err)
		return
	}

	defer f.Close()
	_, err = io.Copy(f, file)
	if err != nil {
		log.Printf("failed to save local image file from uploaded render frame '%d', token '%s': %v", frameNumber, token[0], err)
		return
	}

	_, ok = batchJobs[token[0]]
	if !ok {
		log.Printf("received a rendered frame '%d' for job with token '%s' but couldn't find this as an active job.  master may have been restarted while render was in progress", frameNumber, token[0])
		return
	}

	fmt.Println("finished saving uploaded file")
	job := batchJobs[token[0]]
	job.CompletedFrames = append(job.CompletedFrames, frameNumber)
	job.CompletePercentage = (len(job.CompletedFrames) * 100) / job.NumberOfFrames

	for i, wip := range job.Wips {
		if wip.FrameNumber == frameNumber {
			job.Wips = append(job.Wips[:i], job.Wips[i+1:]...) //Remove it from the Wips
			break
		}
	}

	f.Close()

	src, err := imaging.Open(fmt.Sprintf("public/renderedframes/%s/%s", token[0], handler.Filename))
	if err != nil {
		log.Fatalf("failed to decode uploaded file during thumbnail creation: %v", err)
	}

	dstImageFit := imaging.Fit(src, 240, 135, imaging.Lanczos)

	f, err = os.Create(fmt.Sprintf("public/renderedframes-thumbnail/%s/%s", token[0], handler.Filename))
	if err != nil {
		panic(err)
	}
	defer f.Close()
	png.Encode(f, dstImageFit)

}

// A slave is checking in to tell us he's still working hard.  Also, provides an update on the current job.
func heartbeatHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("received post request for heartbeat")
	err := r.ParseForm()
	if err != nil {
		log.Printf("failed to parse multipart form in heartbeatHandler")
		return
	}
	token, ok := r.Form["token"]
	if !ok {
		log.Printf("couldn't get token for heartbeatHandler")
		return
	}

	status, ok := r.Form["status"]
	if !ok {
		log.Printf("couldn't get status for heartbeatHandler")
	}

	_, ok = r.Form["framenumber"]
	if !ok {
		log.Printf("couldn't get framenumber for heartbeatHandler")
		return
	}

	frameNumber, err := strconv.Atoi(r.Form["framenumber"][0])
	if err != nil {
		log.Printf("failed to convert frame number '%s' to integer: %v", r.Form["framenumber"][0], err)
		return
	}

	_, ok = batchJobs[token[0]]
	if !ok {
		log.Printf("received heartbeat for batch job with token '%s' but it doesn't exist", token[0])
		return
	}

	job := batchJobs[token[0]]
	var progress int
	if strings.Contains(status[0], "View Layer Rendered") {
		progString := strings.Split(status[0], "View Layer Rendered ")[1]
		progString = strings.Split(progString, " Tiles")[0]
		tiles := strings.Split(progString, "/")
		currentTile, err := strconv.Atoi(tiles[0])
		if err != nil {
			log.Printf("couldn't convert current tile '%s' to integer: %v", tiles[0], err)
		}
		endTile, err := strconv.Atoi(tiles[1])
		if err != nil {
			log.Printf("couldn't convert endTile '%s' to integer: %v", tiles[1], err)
		}
		log.Printf("currenttile: %d\nendTile: %d", currentTile, endTile)
		progress = (currentTile * 100) / endTile
	}

	ok = false
	for i, wip := range job.Wips {
		if wip.FrameNumber == frameNumber {
			log.Printf("received heartbeat for batch job '%s' - current status '%s' - current progress '%d'", token[0], status[0], progress)
			ok = true
			job.Wips[i].HeartbeatTimestamp = time.Now()
			job.Wips[i].Output = status[0]
			job.Wips[i].CompletePercentage = progress
			break
		}
	}
	if !ok {
		log.Printf("received heartbeat for frame '%d' for batch job '%s' but the frame is not considered to be in progress?", frameNumber, token[0])
	}
}

// Checks all wips for all jobs for the heartbeat timestamp timeout, and considers the slave down if it's exceeded
// the work then becomes available to be redispatched to the next slave that requests work.
func watchdogTimer() {
	for {
		time.Sleep(time.Minute * 1) // We run this every minute
		for _, job := range batchJobs {
			for i, wip := range job.Wips {
				if time.Since(wip.HeartbeatTimestamp) > time.Minute*1 {
					job.Wips[i].BadHealth = true
				} else {
					job.Wips[i].BadHealth = false
				}
				if time.Since(wip.HeartbeatTimestamp) > time.Minute*5 {
					log.Printf("slave for frame '%d on job '%s' hasn't checked in for over 5 minutes... considering them dead", wip.FrameNumber, job.Token)
					job.RemainingFrames = append([]int{wip.FrameNumber}, job.RemainingFrames...) //Throw it to the front of the queue
					job.Wips = append(job.Wips[:i], job.Wips[i+1:]...)
					break
				}
			}
		}
	}
}

func notFoundHandler(w http.ResponseWriter, r *http.Request) {
	t, err := template.ParseFiles("templates/404.html")
	if err != nil {
		log.Fatalf("could not parse 404.html template file - you know shit's going wrong when...: %v", err)
	}
	t.Execute(w, nil)
}

//A little helper function, removes all whitespace from a string
func stripSpaces(str string) string {
	return strings.Map(func(r rune) rune {
		if unicode.IsSpace(r) {
			// if the character is a space, drop it
			return -1
		}
		// else keep it in the string
		return r
	}, str)
}