github.com/gramework/gramework@v1.8.1-0.20231027140105-82555c9057f5/x/client/wsHandler.go (about)

     1  // Copyright 2017-present Kirill Danshin and Gramework contributors
     2  // Copyright 2019-present Highload LTD (UK CN: 11893420)
     3  //
     4  // Licensed under the Apache License, Version 2.0 (the "License");
     5  // you may not use this file except in compliance with the License.
     6  // You may obtain a copy of the License at
     7  //
     8  //     http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  
    11  package client
    12  
    13  import (
    14  	"time"
    15  
    16  	"github.com/fasthttp/websocket"
    17  	"github.com/gramework/gramework"
    18  )
    19  
    20  // WSHandler returns gramework handler
    21  func (client *Instance) WSHandler() func(*gramework.Context) error {
    22  	up := websocket.FastHTTPUpgrader{}
    23  	return func(ctx *gramework.Context) error {
    24  		if websocket.FastHTTPIsWebSocketUpgrade(ctx.RequestCtx) {
    25  			return up.Upgrade(ctx.RequestCtx, func(conn *websocket.Conn) {
    26  				for {
    27  					v := <-client.watch(ctx)
    28  					err := conn.WriteMessage(websocket.TextMessage, v)
    29  					if err != nil {
    30  						break
    31  					}
    32  				}
    33  			})
    34  		}
    35  
    36  		return client.handleHTTP(ctx)
    37  	}
    38  }
    39  
    40  func (client *Instance) watch(ctx *gramework.Context) chan []byte {
    41  	c := make(chan []byte)
    42  	go func() {
    43  		for {
    44  			api, err := client.nextServer()
    45  			if err != nil {
    46  				ctx.Logger.Errorf("error: %s", err)
    47  				time.Sleep(client.conf.WatcherTickTime)
    48  				continue
    49  			}
    50  			bytes := buffer.Get()
    51  			_, body, err := api.HostClient.Get(bytes.B, api.Addr)
    52  			if err != nil {
    53  				ctx.Logger.Errorf("error while .Do() the request %s", err)
    54  				time.Sleep(client.conf.WatcherTickTime)
    55  				buffer.Put(bytes)
    56  				continue
    57  			}
    58  			c <- body
    59  			time.Sleep(client.conf.WatcherTickTime)
    60  			buffer.Put(bytes)
    61  		}
    62  	}()
    63  
    64  	return c
    65  }