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() } }