telegraf/plugins/inputs/cl_104/upstream.go

160 lines
2.9 KiB
Go

package cl_104
import (
"context"
"encoding/json"
"net/http"
"sync"
"sync/atomic"
"time"
"github.com/gorilla/websocket"
)
var upConnNum int64
func (h *CL104) serveUpstream(res http.ResponseWriter, req *http.Request) {
select {
case <-h.close:
res.WriteHeader(http.StatusGone)
return
default:
if atomic.SwapInt64(&upConnNum, 1) > 0 {
res.WriteHeader(http.StatusConflict)
return
}
}
conn, err := upgrader.Upgrade(res, req, nil)
if err != nil {
h.Log.Error(err)
return
}
defer conn.Close()
stopCtx, stopCancel := context.WithCancel(context.Background())
defer stopCancel()
session := &wsSession{
conn: conn,
ctx: stopCtx,
cancel: stopCancel,
ctrlCh: make(chan wsMsg, 2),
}
if err := h.setUpSessionConn(session); err != nil {
h.Log.Error(err)
return
}
h.startUpWorkers(session)
}
func (h *CL104) startUpWorkers(session *wsSession) {
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
h.monitorUpWrite(session)
}()
wg.Add(1)
go func() {
defer wg.Done()
h.monitorUpRead(session)
}()
wg.Add(1)
go func() {
defer wg.Done()
h.sendPing(session)
}()
wg.Wait()
atomic.SwapInt64(&upConnNum, 0)
}
func (h *CL104) monitorUpWrite(session *wsSession) {
var err error
for {
select {
case <-session.ctx.Done():
return
case ctrl := <-session.ctrlCh:
err = session.conn.SetWriteDeadline(time.Now().Add(time.Duration(h.WriteWait)))
if err != nil {
session.cancel()
return
}
err = session.conn.WriteControl(ctrl.mt, ctrl.data, time.Now().Add(time.Duration(h.WriteWait)))
if err != nil {
session.cancel()
return
}
case msg := <-h.upChan:
err = session.conn.SetWriteDeadline(time.Now().Add(time.Duration(h.WriteWait)))
if err != nil {
session.cancel()
return
}
err = session.conn.WriteMessage(websocket.TextMessage, msg)
if err != nil {
session.cancel()
return
}
}
}
}
func (h *CL104) monitorUpRead(session *wsSession) {
for {
select {
case <-session.ctx.Done():
return
default:
mt, rm, err := session.conn.ReadMessage()
if err != nil {
if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
session.cancel()
return
}
if ce, ok := err.(*websocket.CloseError); ok {
h.Log.Infof("client closed with code:", ce.Code, "text:", ce.Text)
session.cancel()
return
}
h.Log.Error("server read:", err)
session.cancel()
return
}
switch mt {
case websocket.TextMessage:
if err := h.fromUpstream(rm); err != nil {
h.Log.Error(err)
}
default:
h.Log.Info("not text:", string(rm))
}
}
}
}
func (h *CL104) fromUpstream(m []byte) error {
msg := new(msg)
if err := json.Unmarshal(m, msg); err != nil {
return err
}
if msg.TI >= 45 && msg.TI <= 69 || msg.TI >= 100 && msg.TI <= 109 {
select {
case h.clChan <- m:
case <-time.After(time.Second * 5):
h.Log.Error("drop cl msg:", msg)
}
}
return nil
}