From 52d4f4fb385f4fb8940f6cdfa566c9f2d1e24700 Mon Sep 17 00:00:00 2001
From: Starbeamrainbowlabs <sbrl@starbeamrainbowlabs.com>
Date: Thu, 25 Jul 2019 18:56:59 +0100
Subject: [PATCH] Bugfix the new web worker.

Web workers + Tensorflow + Rollup = pain. I think I might blog about
it....
---
 client_src/index.html              |  2 +-
 client_src/js/LayerAI.mjs          | 39 ++++++++++----------
 client_src/js/Worker/AI.worker.mjs | 58 +++++++++++++++++-------------
 client_src/js/Worker/AIWrapper.mjs | 35 +++++++++---------
 package-lock.json                  | 13 +++----
 package.json                       |  2 +-
 rollup.config.js                   | 12 +++----
 7 files changed, 85 insertions(+), 76 deletions(-)

diff --git a/client_src/index.html b/client_src/index.html
index 42804c5..89456e8 100644
--- a/client_src/index.html
+++ b/client_src/index.html
@@ -19,6 +19,6 @@
 		
 		<!---------------->
 		<link rel="stylesheet" href="app.css" />
-		<script src="app.js" charset="utf-8"></script>
+		<script src="app.js" type="module" charset="utf-8"></script>
 	</body>
 </html>
diff --git a/client_src/js/LayerAI.mjs b/client_src/js/LayerAI.mjs
index f6ec9ff..8c68b5b 100644
--- a/client_src/js/LayerAI.mjs
+++ b/client_src/js/LayerAI.mjs
@@ -6,8 +6,6 @@ import chroma from 'chroma-js';
 import GetFromUrl from './Helpers/GetFromUrl.mjs';
 import Config from './ClientConfig.mjs';
 
-import AIWorker from './Worker/AI.worker.mjs';
-
 
 class LayerAI {
 	/**
@@ -39,7 +37,7 @@ class LayerAI {
 	 */
 	constructor(map) {
 		this.map = map;
-		this.worker = new AIWorker();
+		this.worker = new Worker("worker.js"); // We could pass { type: "module" } as the 2nd argument here to allow Rollup code splitting to function as normal, but it's not supported by browsers yet
 		
 		this.map_bounds = null;
 	}
@@ -66,8 +64,11 @@ class LayerAI {
 			// Ask the web worker to set itself up
 			this.worker.postMessage({
 				event: "setup",
-				bounds: this.gateway_bounds,
-				index: this.index
+				setup_info: {
+					bounds: this.map_bounds,
+					index: this.index,
+					Config
+				}
 			});
 		})
 	}
@@ -112,11 +113,11 @@ class LayerAI {
 		
 		// Figure out the bounds of the map we're going to generate
 		this.map_bounds = this.gateway_bounds;
-		map_bounds.north += Config.border.lat;
-		map_bounds.south -= Config.border.lat;
+		this.map_bounds.north += Config.border.lat;
+		this.map_bounds.south -= Config.border.lat;
 		
-		map_bounds.east += Config.border.lng;
-		map_bounds.west -= Config.border.lng;
+		this.map_bounds.east += Config.border.lng;
+		this.map_bounds.west -= Config.border.lng;
 		
 		// Setup the web worker
 		await this.worker_setup();
@@ -134,7 +135,7 @@ class LayerAI {
 	 */
 	async generate_layer() {
 		console.log("[Layer/AI] Rendering map");
-		let map = this.render_map();
+		let map = await this.render_map();
 		console.log("[Layer/AI] Passing to Leaflet");
 		return L.geoJSON(map, {
 			style: (feature) => { return {
@@ -169,17 +170,17 @@ class LayerAI {
 		};
 		
 		for(let lat = this.map_bounds.south; lat < this.map_bounds.north; lat += Config.step.lat) {
-			let next_row = await worker_predict_row(lat);
+			let next_result = await this.worker_predict_row(lat);
 			let lng = this.map_bounds.west;
 			
+			// Keep up with the statistics
+			if(next_result.stats.rssi_max > stats.rssi_max)
+				stats.rssi_max = next_result.stats.rssi_max;
+			if(next_result.stats.rssi_min < stats.rssi_min)
+				stats.rssi_min = next_result.stats.rssi_min;
 			
-			for(let value of next_row) {
-				// Keep up with the statistics
-				if(value > stats.rssi_max)
-					stats.rssi_max = value;
-				if(value < stats.rssi_min)
-					stats.rssi_min = value;
-				
+			
+			for(let value of next_result.result) {
 				// Generate the GeoJSON feature for this cell of the map
 				coverage.push({
 					type: "Feature",
@@ -204,7 +205,7 @@ class LayerAI {
 			}
 		}
 		
-		console.log(stats);
+		console.log("Stats:", stats);
 		
 		return coverage;
 	}
diff --git a/client_src/js/Worker/AI.worker.mjs b/client_src/js/Worker/AI.worker.mjs
index ef60dd2..a7b0155 100644
--- a/client_src/js/Worker/AI.worker.mjs
+++ b/client_src/js/Worker/AI.worker.mjs
@@ -1,26 +1,36 @@
+"use strict";
 
-export default function(self) {
-	let ai_wrapper = new AIWrapper();
-	
-	self.addEventListener("message", async (event) => {
-		console.log(event.data);
-		switch(event.data.event) {
-			case "setup":
-				await ai_wrapper.setup(event.data.setup_info);
-				self.postMessage({
-					"event": "setup-complete"
-				});
-				break;
-			
-			case "predict-row":
-				let message = await ai_wrapper.predict_row(event.data.latitude);
-				message.event = "result";
-				self.postMessage(message);
-				break;
-			
-			case "end":
-				self.close();
-				break;
-		}
-	});
+import AIWrapper from './AIWrapper.mjs';
+
+let ai_wrapper = new AIWrapper();
+
+async function handle_message(event) {
+	console.log(event.data);
+	switch(event.data.event) {
+		case "setup":
+			await ai_wrapper.setup(event.data.setup_info);
+			self.postMessage({
+				"event": "setup-complete"
+			});
+			break;
+		
+		case "predict-row":
+			let message = await ai_wrapper.predict_row(event.data.latitude);
+			message.event = "result";
+			self.postMessage(message);
+			break;
+		
+		case "end":
+			self.close();
+			break;
+		
+		default:
+			throw new Error(`Error: Unknown event '${event.data.event}'`);
+	}
 }
+
+self.addEventListener("message", (event) => {
+	handle_message(event).catch((error) => {
+		console.error(error);
+	});
+});
diff --git a/client_src/js/Worker/AIWrapper.mjs b/client_src/js/Worker/AIWrapper.mjs
index 4b75f66..09cbf79 100644
--- a/client_src/js/Worker/AIWrapper.mjs
+++ b/client_src/js/Worker/AIWrapper.mjs
@@ -4,11 +4,11 @@ import path from 'path';
 
 import {
 	loadLayersModel as tf_loadLayersModel,
-	tensor as tf_tensor
+	tensor as tf_tensor,
+	setBackend as tf_setBackend
 } from '@tensorflow/tfjs';
 
-import { normalise } from '../../common/Math.mjs';
-import Config from '../ClientConfig.mjs';
+import { normalise } from '../../../common/Math.mjs';
 
 class AIWrapper {
 	constructor() {
@@ -16,39 +16,48 @@ class AIWrapper {
 		
 		this.map_bounds = null;
 		this.index = null;
+		this.Config = null;
 		
 		this.gateways = new Map();
 	}
 	
-	async setup({ bounds, index }) {
+	async setup({ bounds, index, Config }) {
 		this.map_bounds = bounds;
 		this.index = index;
+		this.Config = Config;
+		
+		console.log("Loading models");
+		
+		// WebGL isn't available inside WebWorkers yet :-(
+		tf_setBackend("cpu");
 		
 		for(let gateway of this.index.index) {
 			this.gateways.set(
 				gateway.id,
-				await tf_loadLayersModel(`${window.location.href}/${path.dirname(Config.ai_index_file)}/${gateway.id}/model.json`)
+				await tf_loadLayersModel(`${path.dirname(self.location.href)}/${path.dirname(this.Config.ai_index_file)}/${gateway.id}/model.json`)
 			);
 		}
-		
+		console.log("Model setup complete.");
 		this.setup_complete = true;
 	}
 	
 	predict_row(lat) {
-		if(!setup_complete)
+		if(!this.setup_complete)
 			throw new Error("Error: Can't do predictions until the setup is complete.");
 			
-		let results = [],
+		let result = [],
 			stats = {
 				rssi_min: Infinity,
 				rssi_max: -Infinity
 			};
 		
-		for(let lng = this.map_bounds.west; lng < this.map_bounds.east; lng += Config.step.lng) {
+		for(let lng = this.map_bounds.west; lng < this.map_bounds.east; lng += this.Config.step.lng) {
 			let max_predicted_rssi = -Infinity;
 			
 			for(let [, ai] of this.gateways) {
-				let next_prediction = this.predict_value(lat, lng)
+				let next_prediction = ai.predict(
+					tf_tensor([ lat, lng ], [1, 2])
+				).arraySync()[0][0];
 				max_predicted_rssi = Math.max(
 					max_predicted_rssi,
 					next_prediction
@@ -73,12 +82,6 @@ class AIWrapper {
 		
 		return { result, stats };
 	}
-	
-	predict_value(latitude, longitude) {
-		return ai.predict(
-			tf_tensor([ latitude, longitude ], [1, 2])
-		).arraySync()[0][0];
-	}
 }
 
 export default AIWrapper;
diff --git a/package-lock.json b/package-lock.json
index c1e7708..459c36a 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -3892,14 +3892,11 @@
         "terser": "^4.1.0"
       }
     },
-    "rollup-plugin-webworkify": {
-      "version": "0.0.4",
-      "resolved": "https://registry.npmjs.org/rollup-plugin-webworkify/-/rollup-plugin-webworkify-0.0.4.tgz",
-      "integrity": "sha512-GvDPdz7qeakeB26cinVpxPgOBWHx/p3xua/itRXZ/30gGOK5QZeb/iu/9V4M7Zuhwmo2UPR1rViVX8OCqmhi1g==",
-      "dev": true,
-      "requires": {
-        "rollup-pluginutils": "^2"
-      }
+    "rollup-plugin-web-worker-loader": {
+      "version": "0.5.1",
+      "resolved": "https://registry.npmjs.org/rollup-plugin-web-worker-loader/-/rollup-plugin-web-worker-loader-0.5.1.tgz",
+      "integrity": "sha512-mlOtJWXLbvmnVFjtGCS3ykDzawHXJTyA8XbVANO6KBmJNwkoqF0wStvZq7wvf5xtlq9+E8JKy8aptti4PcID6A==",
+      "dev": true
     },
     "rollup-pluginutils": {
       "version": "2.8.1",
diff --git a/package.json b/package.json
index 04d0abb..7108825 100644
--- a/package.json
+++ b/package.json
@@ -44,6 +44,6 @@
 		"rollup-plugin-node-resolve": "^5.2.0",
 		"rollup-plugin-postcss": "^2.0.3",
 		"rollup-plugin-replace": "^2.2.0",
-		"rollup-plugin-webworkify": "0.0.4"
+		"rollup-plugin-web-worker-loader": "^0.5.1"
 	}
 }
diff --git a/rollup.config.js b/rollup.config.js
index 9144143..2bbf945 100644
--- a/rollup.config.js
+++ b/rollup.config.js
@@ -8,7 +8,6 @@ import postcss from 'rollup-plugin-postcss';
 import { terser } from "rollup-plugin-terser";
 import replace from 'rollup-plugin-replace';
 import builtins from '@joseph184/rollup-plugin-node-builtins';
-import webworkify from 'rollup-plugin-webworkify';
 // import json from 'rollup-plugin-json';
 
 import postcss_import from 'postcss-import';
@@ -56,10 +55,6 @@ let plugins = [
 		
 	}),
 	
-	webworkify({
-		pattern: '**/*.worker.mjs'
-	}),
-	
 	postcss({
 		plugins: [
 			postcss_import({}),
@@ -93,9 +88,12 @@ if(process.env.NODE_ENV == "production") {
 }
 
 export default {
-	input: 'client_src/js/index.mjs',
+	input: {
+		"app": 'client_src/js/index.mjs',
+		"worker": 'client_src/js/Worker/AI.worker.mjs'
+	},
 	output: {
-		file: 'app/app.js',
+		dir: 'app/',
 		format: 'esm'
 	},
 	plugins
