github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/internal/http/connection.go (about)

     1  package simplehttp
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net/http"
     7  	"time"
     8  
     9  	"github.com/Asutorufa/yuhaiin/internal/appapi"
    10  	"github.com/Asutorufa/yuhaiin/pkg/log"
    11  	websocket "github.com/Asutorufa/yuhaiin/pkg/net/proxy/websocket/x"
    12  	gs "github.com/Asutorufa/yuhaiin/pkg/protos/statistic/grpc"
    13  	"github.com/Asutorufa/yuhaiin/pkg/utils/relay"
    14  	"google.golang.org/grpc/metadata"
    15  	"google.golang.org/protobuf/types/known/emptypb"
    16  )
    17  
    18  func ConnWebsocket(cc *appapi.Components) func(w http.ResponseWriter, r *http.Request) error {
    19  	return func(w http.ResponseWriter, r *http.Request) error {
    20  		return websocket.ServeHTTP(w, r, func(ctx context.Context, c *websocket.Conn) error {
    21  			defer c.Close()
    22  
    23  			var ticker int
    24  			err := websocket.JSON.Receive(c, &ticker)
    25  			if err != nil {
    26  				return err
    27  			}
    28  			if ticker <= 0 {
    29  				ticker = 2000
    30  			}
    31  
    32  			ctx, cancel := context.WithCancel(ctx)
    33  			defer cancel()
    34  
    35  			cns := &connectionsNotifyServer{ctx, make(chan *gs.NotifyData, 20)}
    36  
    37  			go func() {
    38  				defer cancel()
    39  				err := cc.Connections.Notify(&emptypb.Empty{}, cns)
    40  
    41  				if err != nil {
    42  					log.Warn("connections notify failed", "err", err)
    43  				}
    44  			}()
    45  
    46  			go func() {
    47  				_, _ = relay.Copy(io.Discard, c)
    48  				cancel()
    49  			}()
    50  
    51  			go func() {
    52  				timer := time.NewTicker(time.Duration(ticker) * time.Millisecond)
    53  				defer timer.Stop()
    54  
    55  				send := func() {
    56  					total, err := cc.Connections.Total(ctx, &emptypb.Empty{})
    57  					if err == nil {
    58  						_ = cns.Send(&gs.NotifyData{
    59  							Data: &gs.NotifyData_TotalFlow{
    60  								TotalFlow: total,
    61  							},
    62  						},
    63  						)
    64  					}
    65  				}
    66  
    67  				send()
    68  
    69  				for {
    70  					select {
    71  					case <-timer.C:
    72  						send()
    73  
    74  					case <-ctx.Done():
    75  						return
    76  					}
    77  				}
    78  			}()
    79  
    80  			for {
    81  				select {
    82  				case <-ctx.Done():
    83  					return nil
    84  
    85  				case m := <-cns.msgChan:
    86  					if err = websocket.PROTO.Send(c, m); err != nil {
    87  						return err
    88  					}
    89  				}
    90  			}
    91  		})
    92  	}
    93  }
    94  
    95  type connectionsNotifyServer struct {
    96  	ctx     context.Context
    97  	msgChan chan *gs.NotifyData
    98  }
    99  
   100  func (x *connectionsNotifyServer) Send(m *gs.NotifyData) error {
   101  	select {
   102  	case <-x.ctx.Done():
   103  		return x.ctx.Err()
   104  	case x.msgChan <- m:
   105  	}
   106  
   107  	return nil
   108  }
   109  
   110  func (x *connectionsNotifyServer) Context() context.Context     { return x.ctx }
   111  func (x *connectionsNotifyServer) SetHeader(metadata.MD) error  { return nil }
   112  func (x *connectionsNotifyServer) SendHeader(metadata.MD) error { return nil }
   113  func (x *connectionsNotifyServer) SetTrailer(metadata.MD)       {}
   114  func (x *connectionsNotifyServer) SendMsg(m any) error          { return nil }
   115  func (x *connectionsNotifyServer) RecvMsg(m any) error          { return nil }