sfcs/client/Start.go

package client

import (
	"log"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"time"

	"github.com/gorilla/websocket"
	"github.com/inhies/go-bytesize"
	"github.com/mackerelio/go-osstat/memory"
	"github.com/shirou/gopsutil/v4/cpu"
)

func (client *Client) Start() {
	// Setup the interrupts
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	u := url.URL{Scheme: "ws", Host: client.ServerAddress + ":" + strconv.Itoa(client.ServerPort), Path: "/ws"}
	log.Printf("Connecting to %s", u.String())

	// sigh
	var err error
	client.conn, _, err = websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		log.Fatal("dial:", err)
	}
	defer client.conn.Close()

	done := make(chan struct{})

	// Setup updater logic
	updater := make(chan string)
	go client.statusUpdater(updater)

	go func() {
		defer close(done)
		for {
			_, message, err := client.conn.ReadMessage()
			if err != nil {
				log.Println("read:", err)
				return
			}
			log.Printf("recv: %s", message)
		}
	}()

	for {
		select {
		case <-done:
			return
		case update := <-updater: // Dumping this into a var saves a second. Weird!
			// Dump the message towards the server.
			err = client.conn.WriteMessage(websocket.TextMessage, []byte(update))
			if err != nil {
				log.Println("write:", err)
				return
			}
		case <-interrupt:
			log.Println("interrupt")

			// Cleanly close the connection by sending a close message and then
			// waiting (with timeout) for the server to close the connection.
			err := client.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
			if err != nil {
				log.Println("write close:", err)
				return
			}
			select {
			case <-done:
			case <-time.After(time.Second):
			}
			return
		}
	}

}

func (client *Client) statusUpdater(updater chan string) {
	for {
		// Fetch CPU usage in percentages
		cpustats, err := cpu.Percent(time.Second, false)
		if err != nil {
			log.Println("Err:", err)
		}

		// Fetch memory stats in bytes
		mem, err := memory.Get()
		if err != nil {
			log.Println("Err:", err)
		}

		// Use the ByteSize package to allow for memory calculations
		b := bytesize.New(float64(mem.Used))

		// Dump the message towards the server.
		updater <- "update;" + strconv.Itoa(int(cpustats[0])) + ";" + b.String() + ";" + strconv.Itoa(client.taskId)
	}
}