ran-proc/httppm/client.go

package httppm

import (
	"context"
	"fmt"
	"github.com/gorilla/websocket"
	"net/http"
	"os"
	"strconv"
	"time"
)

type HPMClientLogFunc func(msg string)

type HPMClient struct {
	hpmURL             string
	checkAliveDuration time.Duration
	running            bool
	closeCncl          context.CancelFunc
	logFunc            HPMClientLogFunc
}

func NewHPMClient(hpmURL string, checkAliveDuration time.Duration) *HPMClient {
	c := &HPMClient{
		hpmURL:             hpmURL,
		checkAliveDuration: checkAliveDuration,
		running:            false,
		logFunc:            func(msg string) {},
	}
	return c
}

func (c *HPMClient) SetLogFunc(fn HPMClientLogFunc) {
	c.logFunc = fn
}

func (c *HPMClient) Run() error {
	pid := os.Getpid()
	hdr := http.Header{}
	hdr.Add("pid", strconv.Itoa(pid))
	wsconn, _, err := websocket.DefaultDialer.Dial(c.hpmURL, hdr)
	if err != nil {
		return fmt.Errorf("HPM Client WS Dial Error: %v", err)
	}
	c.logFunc("hpm client websocket connected.")
	hbTimer := time.NewTicker(c.checkAliveDuration)
	ctx, cncl := context.WithCancel(context.Background())
	c.closeCncl = cncl
	c.running = true
	defer hbTimer.Stop()
	go func() {
		for {
			_, _, rerr := wsconn.ReadMessage()
			if rerr != nil {
				c.logFunc(fmt.Sprint("hpm stop by ws recv fail, detail: ", rerr))
				_ = wsconn.Close()
				cncl()
				return
			}
		}
	}()
MLoop:
	for {
		select {
		case <-hbTimer.C:
			xerr := wsconn.WriteMessage(websocket.PingMessage, nil)
			if xerr != nil {
				c.logFunc(fmt.Sprint("hpm stop by ws send fail, detail: ", xerr))
				cncl()
				continue MLoop
			}
		case <-ctx.Done():
			{
				c.logFunc("hpm stopped.")
				return nil
			}
		}
	}
	return nil
}

func (c *HPMClient) Stop(xerr error) {
	if c.running {
		c.logFunc("hpm stop by other.")
		c.closeCncl()
	}
}