diff --git a/wshandler/wshandler.go b/wshandler/wshandler.go index 2e56fe1..d2860ea 100644 --- a/wshandler/wshandler.go +++ b/wshandler/wshandler.go @@ -188,6 +188,7 @@ type websocketHandlerBase struct { type WebsocketHandler struct { WebsocketApiBroker websocketHandlerBase + opt *WebsocketHandlerOption } type wsConfig struct { @@ -205,7 +206,28 @@ func init() { gob.Register([]any{}) } -func NewWebsocketHandler(consumer session.Consumer, redisUrl string) (*WebsocketHandler, error) { +type WebsocketHandlerOption struct { + ReadBufferSize, WriteBufferSize int +} + +func Option() *WebsocketHandlerOption { + return &WebsocketHandlerOption{ + ReadBufferSize: 0, + WriteBufferSize: 0, + } +} + +func (opt *WebsocketHandlerOption) SetReadBufferSize(size int) *WebsocketHandlerOption { + opt.ReadBufferSize = size + return opt +} + +func (opt *WebsocketHandlerOption) SetWriteBufferSize(size int) *WebsocketHandlerOption { + opt.WriteBufferSize = size + return opt +} + +func NewWebsocketHandler(consumer session.Consumer, redisUrl string, opts ...*WebsocketHandlerOption) (*WebsocketHandler, error) { var config wsConfig if err := gocommon.LoadConfig(&config); err != nil { return nil, err @@ -240,6 +262,14 @@ func NewWebsocketHandler(consumer session.Consumer, redisUrl string) (*Websocket } }() + var opt *WebsocketHandlerOption + if len(opts) > 0 { + // TODO : opts merge + opt = opts[0] + } else { + opt = Option() + } + ws := &WebsocketHandler{ websocketHandlerBase: websocketHandlerBase{ redisMsgChanName: fmt.Sprintf("_wsh_msg_%d", redisSync.Options().DB), @@ -252,6 +282,7 @@ func NewWebsocketHandler(consumer session.Consumer, redisUrl string) (*Websocket sendMsgChan: sendchan, sessionConsumer: consumer, }, + opt: opt, } consumer.RegisterOnSessionInvalidated(ws.onSessionInvalidated) return ws, nil @@ -665,7 +696,10 @@ func (ws *WebsocketHandler) upgrade_nosession(w http.ResponseWriter, r *http.Req return } - var upgrader = websocket.Upgrader{} // use default options + var upgrader = websocket.Upgrader{ + ReadBufferSize: ws.opt.ReadBufferSize, + WriteBufferSize: ws.opt.WriteBufferSize, + } // use default options conn, err := upgrader.Upgrade(w, r, nil) if err != nil { w.WriteHeader(http.StatusInternalServerError)