From d40497aca59040695e05ecbf840d0da8bb29afc1 Mon Sep 17 00:00:00 2001
From: Jake McDermott <yo@jakemcdermott.me>
Date: Thu, 14 May 2020 12:46:49 -0400
Subject: [PATCH] Add basic output updates with websockets

---
 awx/ui_next/README.md                         |  2 +
 awx/ui_next/package.json                      |  3 +-
 .../src/screens/Job/JobOutput/JobOutput.jsx   | 61 +++++++++++++++++++
 awx/ui_next/src/setupProxy.js                 | 11 +---
 4 files changed, 68 insertions(+), 9 deletions(-)

diff --git a/awx/ui_next/README.md b/awx/ui_next/README.md
index 5ab552b401..d473bcf79c 100644
--- a/awx/ui_next/README.md
+++ b/awx/ui_next/README.md
@@ -20,6 +20,8 @@ you'll need to update your django settings and use the `TARGET_HOST` and `TARGET
 echo "CSRF_TRUSTED_ORIGINS = ['awx.local:8043']" >> /awx/settings/development.py
 TARGET_HOST='awx.local:8043' TARGET_PORT=8043 npm --prefix awx/ui_next start
 ```
+**Note:** When using an external server, you must also manually update the `proxy` field in `package.json`
+to point to the new websocket url.
 
 ## Testing
 ```shell
diff --git a/awx/ui_next/package.json b/awx/ui_next/package.json
index a8226dea75..64af039038 100644
--- a/awx/ui_next/package.json
+++ b/awx/ui_next/package.json
@@ -87,5 +87,6 @@
       "<rootDir>/src/locales",
       "index.js"
     ]
-  }
+  },
+  "proxy": "https://localhost:8043/websocket"
 }
diff --git a/awx/ui_next/src/screens/Job/JobOutput/JobOutput.jsx b/awx/ui_next/src/screens/Job/JobOutput/JobOutput.jsx
index 724fd845b2..4e72d0014c 100644
--- a/awx/ui_next/src/screens/Job/JobOutput/JobOutput.jsx
+++ b/awx/ui_next/src/screens/Job/JobOutput/JobOutput.jsx
@@ -63,6 +63,44 @@ const OutputFooter = styled.div`
   flex: 1;
 `;
 
+let ws;
+function connectJobSocket({ type, id }, onMessage) {
+  ws = new WebSocket(`wss://${window.location.host}/websocket/`);
+
+  ws.onopen = () => {
+    const xrftoken = `; ${document.cookie}`
+      .split('; csrftoken=')
+      .pop()
+      .split(';')
+      .shift();
+    const eventGroup = `${type}_events`;
+    ws.send(
+      JSON.stringify({
+        xrftoken,
+        groups: { jobs: ['summary', 'status_changed'], [eventGroup]: [id] },
+      })
+    );
+  };
+
+  ws.onmessage = e => {
+    onMessage(JSON.parse(e.data));
+  };
+
+  ws.onclose = e => {
+    // eslint-disable-next-line no-console
+    console.debug('Socket closed. Reconnecting...', e);
+    setTimeout(() => {
+      connectJobSocket({ type, id }, onMessage);
+    }, 1000);
+  };
+
+  ws.onerror = err => {
+    // eslint-disable-next-line no-console
+    console.debug('Socket error: ', err, 'Disconnecting...');
+    ws.close();
+  };
+}
+
 function range(low, high) {
   const numbers = [];
   for (let n = low; n <= high; n++) {
@@ -105,11 +143,22 @@ class JobOutput extends Component {
     this.isRowLoaded = this.isRowLoaded.bind(this);
     this.loadMoreRows = this.loadMoreRows.bind(this);
     this.scrollToRow = this.scrollToRow.bind(this);
+    this.monitorJobSocketCounter = this.monitorJobSocketCounter.bind(this);
   }
 
   componentDidMount() {
+    const { job } = this.props;
     this._isMounted = true;
     this.loadJobEvents();
+
+    connectJobSocket(job, data => {
+      if (data.counter && data.counter > this.jobSocketCounter) {
+        this.jobSocketCounter = data.counter;
+      } else if (data.final_counter && data.unified_job_id === job.id) {
+        this.jobSocketCounter = data.final_counter;
+      }
+    });
+    this.interval = setInterval(() => this.monitorJobSocketCounter(), 5000);
   }
 
   componentDidUpdate(prevProps, prevState) {
@@ -131,9 +180,21 @@ class JobOutput extends Component {
   }
 
   componentWillUnmount() {
+    if (ws) {
+      ws.close();
+    }
+    clearInterval(this.interval);
     this._isMounted = false;
   }
 
+  monitorJobSocketCounter() {
+    const { remoteRowCount } = this.state;
+    if (this.jobSocketCounter >= remoteRowCount) {
+      this._isMounted &&
+        this.setState({ remoteRowCount: this.jobSocketCounter + 1 });
+    }
+  }
+
   async loadJobEvents() {
     const { job, type } = this.props;
 
diff --git a/awx/ui_next/src/setupProxy.js b/awx/ui_next/src/setupProxy.js
index be63698746..cce025003f 100644
--- a/awx/ui_next/src/setupProxy.js
+++ b/awx/ui_next/src/setupProxy.js
@@ -4,6 +4,9 @@ const TARGET_PORT = process.env.TARGET_PORT || 8043;
 const TARGET_HOST = process.env.TARGET_HOST || 'localhost';
 const TARGET = `https://${TARGET_HOST}:${TARGET_PORT}`;
 
+// Note: The websocket proxy is configured
+// manually using the 'proxy' field in package.json
+
 module.exports = app => {
   app.use(
     '/api/login/',
@@ -29,12 +32,4 @@ module.exports = app => {
         req.originalUrl.includes('login'),
     })
   );
-  app.use(
-    '/websocket',
-    createProxyMiddleware({
-      target: TARGET,
-      secure: false,
-      ws: true,
-    })
-  );
 };