wslog.go

Documentation: github.com/cybertec-postgresql/pgwatch/v3/internal/webserver

     1  package webserver
     2  
     3  import (
     4  	"net/http"
     5  	"time"
     6  
     7  	"github.com/cybertec-postgresql/pgwatch/v3/internal/log"
     8  	"github.com/gorilla/websocket"
     9  )
    10  
    11  const (
    12  	// Time allowed to write the file to the client.
    13  	writeWait = 10 * time.Second
    14  
    15  	// Time allowed to read the next pong message from the client.
    16  	pongWait = 60 * time.Second
    17  
    18  	// Send pings to client with this period. Must be less than pongWait.
    19  	pingPeriod = (pongWait * 9) / 10
    20  )
    21  
    22  var (
    23  	upgrader = websocket.Upgrader{
    24  		ReadBufferSize:  1024,
    25  		WriteBufferSize: 1024,
    26  	}
    27  )
    28  
    29  func reader(ws *websocket.Conn) {
    30  	defer ws.Close()
    31  	ws.SetReadLimit(512)
    32  	err := ws.SetReadDeadline(time.Now().Add(pongWait))
    33  	if err != nil {
    34  		return
    35  	}
    36  	ws.SetPongHandler(func(string) error { return ws.SetReadDeadline(time.Now().Add(pongWait)) })
    37  	for {
    38  		_, _, err = ws.ReadMessage()
    39  		if err != nil {
    40  			break
    41  		}
    42  	}
    43  }
    44  
    45  func writer(ws *websocket.Conn, l log.LoggerHookerIface) {
    46  	pingTicker := time.NewTicker(pingPeriod)
    47  	defer func() {
    48  		pingTicker.Stop()
    49  		ws.Close()
    50  	}()
    51  	msgChan := make(log.MessageChanType)
    52  	l.AddSubscriber(msgChan)
    53  	defer l.RemoveSubscriber(msgChan)
    54  	for {
    55  		select {
    56  		case msg := <-msgChan:
    57  			if ws.SetWriteDeadline(time.Now().Add(writeWait)) != nil ||
    58  				ws.WriteMessage(websocket.TextMessage, []byte(msg)) != nil {
    59  				return
    60  			}
    61  		case <-pingTicker.C:
    62  			if ws.SetWriteDeadline(time.Now().Add(writeWait)) != nil ||
    63  				ws.WriteMessage(websocket.PingMessage, []byte{}) != nil {
    64  				return
    65  			}
    66  		}
    67  	}
    68  }
    69  
    70  func (Server *WebUIServer) serveWsLog(w http.ResponseWriter, r *http.Request) {
    71  	ws, err := upgrader.Upgrade(w, r, nil)
    72  	if err != nil {
    73  		Server.l.Error(err)
    74  		return
    75  	}
    76  	Server.l.WithField("reguest", r.URL.String()).Debugf("established websocket connection")
    77  	l, ok := Server.l.(log.LoggerHookerIface)
    78  	if !ok {
    79  		Server.l.Error("cannot cast WebUIServer.l to log.LoggerHookerIface")
    80  		return
    81  	}
    82  	go writer(ws, l)
    83  	reader(ws)
    84  }
    85