github.com/clubpay/ronykit/kit@v0.14.4-0.20240515065620-d0dace45cbc7/stub/stub_ws_option.go (about)

     1  package stub
     2  
     3  import (
     4  	"compress/flate"
     5  	"net/http"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/clubpay/ronykit/kit"
    10  	"github.com/fasthttp/websocket"
    11  )
    12  
    13  type Dialer = websocket.Dialer
    14  
    15  type PreDialHandler func(d *Dialer)
    16  
    17  type OnConnectHandler func(ctx *WebsocketCtx)
    18  
    19  type WebsocketOption func(cfg *wsConfig)
    20  
    21  type wsConfig struct {
    22  	predicateKey    string
    23  	rpcInFactory    kit.IncomingRPCFactory
    24  	rpcOutFactory   kit.OutgoingRPCFactory
    25  	ratelimitChan   chan struct{}
    26  	handlersWG      sync.WaitGroup
    27  	handlers        map[string]RPCContainerHandler
    28  	defaultHandler  RPCContainerHandler
    29  	tracePropagator kit.TracePropagator
    30  
    31  	autoReconnect bool
    32  	dialerBuilder func() *websocket.Dialer
    33  	upgradeHdr    http.Header
    34  	pingTime      time.Duration
    35  	dialTimeout   time.Duration
    36  	writeTimeout  time.Duration
    37  	compressLevel int
    38  
    39  	preDial    func(d *websocket.Dialer)
    40  	onConnect  OnConnectHandler
    41  	preflights []RPCPreflightHandler
    42  
    43  	panicRecoverFunc func(err any)
    44  }
    45  
    46  func WithUpgradeHeader(key string, values ...string) WebsocketOption {
    47  	return func(cfg *wsConfig) {
    48  		if cfg.upgradeHdr == nil {
    49  			cfg.upgradeHdr = http.Header{}
    50  		}
    51  		cfg.upgradeHdr[key] = values
    52  	}
    53  }
    54  
    55  func WithCustomDialerBuilder(b func() *websocket.Dialer) WebsocketOption {
    56  	return func(cfg *wsConfig) {
    57  		cfg.dialerBuilder = b
    58  	}
    59  }
    60  
    61  func WithDefaultHandler(h RPCContainerHandler) WebsocketOption {
    62  	return func(cfg *wsConfig) {
    63  		cfg.defaultHandler = h
    64  	}
    65  }
    66  
    67  func WithHandler(predicate string, h RPCContainerHandler) WebsocketOption {
    68  	return func(cfg *wsConfig) {
    69  		if cfg.handlers == nil {
    70  			cfg.handlers = map[string]RPCContainerHandler{}
    71  		}
    72  		cfg.handlers[predicate] = h
    73  	}
    74  }
    75  
    76  func WithCustomRPC(in kit.IncomingRPCFactory, out kit.OutgoingRPCFactory) WebsocketOption {
    77  	return func(cfg *wsConfig) {
    78  		cfg.rpcInFactory = in
    79  		cfg.rpcOutFactory = out
    80  	}
    81  }
    82  
    83  func WithOnConnectHandler(f OnConnectHandler) WebsocketOption {
    84  	return func(cfg *wsConfig) {
    85  		cfg.onConnect = f
    86  	}
    87  }
    88  
    89  func WithPreDialHandler(f PreDialHandler) WebsocketOption {
    90  	return func(cfg *wsConfig) {
    91  		cfg.preDial = f
    92  	}
    93  }
    94  
    95  func WithPredicateKey(key string) WebsocketOption {
    96  	return func(cfg *wsConfig) {
    97  		cfg.predicateKey = key
    98  	}
    99  }
   100  
   101  func WithAutoReconnect(b bool) WebsocketOption {
   102  	return func(cfg *wsConfig) {
   103  		cfg.autoReconnect = b
   104  	}
   105  }
   106  
   107  func WithPingTime(t time.Duration) WebsocketOption {
   108  	return func(cfg *wsConfig) {
   109  		cfg.pingTime = t
   110  	}
   111  }
   112  
   113  func WithConcurrency(n int) WebsocketOption {
   114  	return func(cfg *wsConfig) {
   115  		cfg.ratelimitChan = make(chan struct{}, n)
   116  	}
   117  }
   118  
   119  func WithRecoverPanic(f func(err any)) WebsocketOption {
   120  	return func(cfg *wsConfig) {
   121  		cfg.panicRecoverFunc = f
   122  	}
   123  }
   124  
   125  // WithPreflightRPC register one or many handlers to run in sequence before
   126  // actually making requests.
   127  func WithPreflightRPC(h ...RPCPreflightHandler) WebsocketOption {
   128  	return func(cfg *wsConfig) {
   129  		cfg.preflights = append(cfg.preflights[:0], h...)
   130  	}
   131  }
   132  
   133  type CompressionLevel int
   134  
   135  const (
   136  	CompressionBestSpeed       CompressionLevel = flate.BestSpeed
   137  	CompressionBestCompression CompressionLevel = flate.BestCompression
   138  )
   139  
   140  func WithCompression(c CompressionLevel) WebsocketOption {
   141  	return func(cfg *wsConfig) {
   142  		switch c {
   143  		default:
   144  		case CompressionBestSpeed, CompressionBestCompression:
   145  			cfg.compressLevel = int(c)
   146  		}
   147  	}
   148  }