├── .gitignore ├── stat ├── counter │ ├── counter.go │ ├── abs │ │ └── abs.go │ ├── avg │ │ └── avg.go │ └── per │ │ └── per.go ├── results │ ├── results.go │ ├── sorter │ │ └── sorter.go │ └── report │ │ └── report.go └── stat.go ├── cli ├── color │ └── color.go ├── cli.go └── input │ └── input.go ├── server ├── flag.go ├── responder.go └── server.go ├── util ├── util.go └── headers │ └── headers.go ├── lua ├── script │ └── script.go ├── util │ └── util.go ├── mod │ ├── ws │ │ ├── server.go │ │ ├── ws.go │ │ └── conn.go │ ├── runtime │ │ └── runtime.go │ ├── time │ │ └── time.go │ ├── mod.go │ └── stat │ │ └── stat.go └── lua.go ├── scripts ├── server.lua ├── load.lua └── client.lua ├── gws.go ├── display ├── display_test.go └── display.go ├── ws ├── connection.go ├── server.go └── ws.go ├── README.md ├── ev ├── timer │ └── timer.go ├── ws │ ├── server.go │ └── client.go └── ev.go ├── client └── client.go ├── bufio └── writer.go └── config └── config.go /.gitignore: -------------------------------------------------------------------------------- 1 | *.iml 2 | .idea 3 | *.json -------------------------------------------------------------------------------- /stat/counter/counter.go: -------------------------------------------------------------------------------- 1 | package counter 2 | -------------------------------------------------------------------------------- /stat/results/results.go: -------------------------------------------------------------------------------- 1 | package results 2 | 3 | type Result struct { 4 | Name string 5 | Kind string 6 | Value float64 7 | Tags map[string]string 8 | Meta map[string]interface{} 9 | } 10 | -------------------------------------------------------------------------------- /cli/color/color.go: -------------------------------------------------------------------------------- 1 | package color 2 | 3 | import ( 4 | "github.com/fatih/color" 5 | ) 6 | 7 | var ( 8 | Red = color.New(color.FgRed).SprintFunc() 9 | Magenta = color.New(color.FgMagenta).SprintFunc() 10 | Green = color.New(color.FgGreen).SprintFunc() 11 | Cyan = color.New(color.FgCyan).SprintFunc() 12 | Yellow = color.New(color.FgYellow).SprintFunc() 13 | ) 14 | -------------------------------------------------------------------------------- /server/flag.go: -------------------------------------------------------------------------------- 1 | package server 2 | 3 | import "fmt" 4 | 5 | type ResponderFlag struct { 6 | value string 7 | expect []string 8 | } 9 | 10 | func (r *ResponderFlag) Set(s string) error { 11 | for _, e := range r.expect { 12 | if e == s { 13 | r.value = s 14 | return nil 15 | } 16 | } 17 | 18 | return fmt.Errorf("expecting one of %s", r.expect) 19 | } 20 | func (r ResponderFlag) String() string { 21 | return r.value 22 | } 23 | func (r ResponderFlag) Get() interface{} { 24 | return r.value 25 | } 26 | -------------------------------------------------------------------------------- /util/util.go: -------------------------------------------------------------------------------- 1 | package util 2 | 3 | import ( 4 | "fmt" 5 | "net/http" 6 | "net/http/httputil" 7 | ) 8 | 9 | func DumpRequestResponse(resp *http.Response) ([]byte, []byte, error) { 10 | if resp == nil { 11 | return nil, nil, fmt.Errorf("nil response") 12 | } 13 | 14 | if resp.Request == nil { 15 | return nil, nil, fmt.Errorf("nil request") 16 | } 17 | 18 | req, err := httputil.DumpRequest(resp.Request, true) 19 | if err != nil { 20 | return nil, nil, err 21 | } 22 | 23 | res, err := httputil.DumpResponse(resp, true) 24 | if err != nil { 25 | return req, nil, err 26 | } 27 | 28 | return req, res, nil 29 | } 30 | -------------------------------------------------------------------------------- /stat/counter/abs/abs.go: -------------------------------------------------------------------------------- 1 | package abs 2 | 3 | import ( 4 | "sync" 5 | ) 6 | 7 | func New() *Abs { 8 | return &Abs{} 9 | } 10 | 11 | type Abs struct { 12 | mu sync.Mutex 13 | value float64 14 | } 15 | 16 | func (a *Abs) Add(v float64) { 17 | a.mu.Lock() 18 | { 19 | a.value += v 20 | } 21 | a.mu.Unlock() 22 | } 23 | 24 | func (a *Abs) Reset() { 25 | a.mu.Lock() 26 | { 27 | a.value = 0 28 | } 29 | a.mu.Unlock() 30 | } 31 | 32 | func (a *Abs) Flush() (result float64) { 33 | a.mu.Lock() 34 | { 35 | result = a.value 36 | } 37 | a.mu.Unlock() 38 | return 39 | } 40 | 41 | func (a *Abs) Kind() string { 42 | return "abs" 43 | } 44 | -------------------------------------------------------------------------------- /stat/counter/avg/avg.go: -------------------------------------------------------------------------------- 1 | package avg 2 | 3 | import ( 4 | "sync" 5 | ) 6 | 7 | func New() *Avg { 8 | return &Avg{} 9 | } 10 | 11 | type Avg struct { 12 | mu sync.Mutex 13 | count float64 14 | value float64 15 | } 16 | 17 | func (a *Avg) Add(v float64) { 18 | a.mu.Lock() 19 | { 20 | a.count++ 21 | a.value += v 22 | } 23 | a.mu.Unlock() 24 | } 25 | 26 | func (a *Avg) Reset() { 27 | a.mu.Lock() 28 | { 29 | a.count = 0 30 | a.value = 0 31 | } 32 | a.mu.Unlock() 33 | } 34 | 35 | func (a *Avg) Flush() (result float64) { 36 | a.mu.Lock() 37 | { 38 | if a.count == 0 { 39 | result = 0 40 | } else { 41 | result = a.value / a.count 42 | } 43 | } 44 | a.mu.Unlock() 45 | return 46 | } 47 | 48 | func (a *Avg) Kind() string { 49 | return "avg" 50 | } 51 | -------------------------------------------------------------------------------- /stat/counter/per/per.go: -------------------------------------------------------------------------------- 1 | package per 2 | 3 | import ( 4 | "sync" 5 | "time" 6 | ) 7 | 8 | func New(d time.Duration) *Per { 9 | return &Per{ 10 | interval: d, 11 | stamp: time.Now(), 12 | } 13 | } 14 | 15 | type Per struct { 16 | mu sync.Mutex 17 | value float64 18 | interval time.Duration 19 | stamp time.Time 20 | } 21 | 22 | func (p *Per) Add(v float64) { 23 | p.mu.Lock() 24 | { 25 | p.value += v 26 | } 27 | p.mu.Unlock() 28 | } 29 | 30 | func (p *Per) Reset() { 31 | p.mu.Lock() 32 | { 33 | p.value = 0 34 | p.stamp = time.Now() 35 | } 36 | p.mu.Unlock() 37 | } 38 | 39 | func (p *Per) Flush() (result float64) { 40 | p.mu.Lock() 41 | { 42 | k := float64(time.Since(p.stamp)) / float64(p.interval) 43 | if k == 0 { 44 | result = 0 45 | } else { 46 | result = p.value / k 47 | } 48 | } 49 | p.mu.Unlock() 50 | return 51 | } 52 | 53 | func (p *Per) Kind() string { 54 | return "per " + p.interval.String() 55 | } 56 | -------------------------------------------------------------------------------- /util/headers/headers.go: -------------------------------------------------------------------------------- 1 | package headers 2 | 3 | import ( 4 | "errors" 5 | "net/http" 6 | "strings" 7 | ) 8 | 9 | var ErrorMalformedHeaderString = errors.New("malformed headers") 10 | 11 | const Separator = ";" 12 | const AssignmentOperator = ':' 13 | 14 | func ParseOne(s string) (key, value string, err error) { 15 | for i := 0; i < len(s); i++ { 16 | if s[i] == AssignmentOperator { 17 | key, value = strings.TrimSpace(s[:i]), strings.TrimSpace(s[i+1:]) 18 | return 19 | } 20 | } 21 | return s, "", ErrorMalformedHeaderString 22 | } 23 | 24 | func Parse(s string) (h http.Header, err error) { 25 | h = make(http.Header) 26 | if s == "" { 27 | return 28 | } 29 | for _, pair := range strings.Split(s, Separator) { 30 | i := strings.Index(pair, string(AssignmentOperator)) 31 | if i == -1 { 32 | err = ErrorMalformedHeaderString 33 | return 34 | } 35 | 36 | h.Add(strings.TrimSpace(pair[:i]), strings.TrimSpace(pair[i+1:])) 37 | } 38 | return 39 | } 40 | -------------------------------------------------------------------------------- /lua/script/script.go: -------------------------------------------------------------------------------- 1 | package script 2 | 3 | import ( 4 | "github.com/gobwas/gws/lua/mod" 5 | "github.com/yuin/gopher-lua" 6 | "io" 7 | ) 8 | 9 | type Script struct { 10 | luaState *lua.LState 11 | } 12 | 13 | func New() *Script { 14 | return &Script{ 15 | luaState: lua.NewState(), 16 | } 17 | } 18 | 19 | func (s *Script) Preload(name string, m mod.Module) { 20 | s.luaState.PreloadModule(name, m.Exports()) 21 | } 22 | 23 | func (s *Script) Do(code string) error { 24 | return s.luaState.DoString(code) 25 | } 26 | 27 | func (s *Script) Shutdown() { 28 | s.luaState.Close() 29 | } 30 | 31 | func (s *Script) HijackOutput(w io.Writer) { 32 | s.luaState.SetGlobal("print", s.luaState.NewFunction(func(L *lua.LState) int { 33 | var buf []byte 34 | for i := 1; ; i++ { 35 | def := L.Get(i) 36 | if _, ok := def.(*lua.LNilType); ok { 37 | break 38 | } 39 | buf = append(buf, []byte(def.String())...) 40 | } 41 | w.Write(append(buf, '\n')) 42 | return 0 43 | })) 44 | } 45 | -------------------------------------------------------------------------------- /server/responder.go: -------------------------------------------------------------------------------- 1 | package server 2 | 3 | import ( 4 | "github.com/chzyer/readline" 5 | "github.com/gobwas/gws/cli/color" 6 | "github.com/gobwas/gws/cli/input" 7 | "github.com/gobwas/gws/ws" 8 | ) 9 | 10 | func DevNullResponder(t ws.Kind, msg []byte) ([]byte, error) { 11 | return nil, nil 12 | } 13 | 14 | func EchoResponder(t ws.Kind, msg []byte) ([]byte, error) { 15 | return msg, nil 16 | } 17 | 18 | func MirrorResponder(t ws.Kind, msg []byte) (r []byte, err error) { 19 | if t != ws.TextMessage { 20 | return 21 | } 22 | 23 | resp := []rune(string(msg)) 24 | for i, l := 0, len(resp)-1; i < len(resp)/2; i, l = i+1, l-1 { 25 | resp[i], resp[l] = resp[l], resp[i] 26 | } 27 | 28 | return []byte(string(resp)), nil 29 | } 30 | 31 | func PromptResponder(t ws.Kind, msg []byte) (r []byte, err error) { 32 | r, err = input.ReadLine(&readline.Config{ 33 | Prompt: color.Green("> "), 34 | HistoryFile: "/tmp/gws_readline_server.tmp", 35 | }) 36 | if err == readline.ErrInterrupt { 37 | return nil, nil 38 | } 39 | 40 | return 41 | 42 | // fmt.Printf("\033[F") 43 | } 44 | -------------------------------------------------------------------------------- /cli/cli.go: -------------------------------------------------------------------------------- 1 | package cli 2 | 3 | import ( 4 | "fmt" 5 | "strings" 6 | ) 7 | 8 | type prefix string 9 | 10 | const ( 11 | PrefixEmpty = " " 12 | PrefixInput = ">" 13 | PrefixIncoming = "<" 14 | PrefixInfo = "–" 15 | PrefixTheEnd = "\u2020" 16 | PrefixBlockStart = "(" 17 | PrefixBlockEnd = ")" 18 | PrefixRaw = "r" 19 | ) 20 | 21 | const ( 22 | PaddingLeft = " " 23 | ) 24 | 25 | func Printf(prefix prefix, format string, c ...interface{}) { 26 | var ( 27 | padLeft, end string 28 | ) 29 | 30 | padLeft = PaddingLeft 31 | end = fmt.Sprintf(" \n%s%s ", padLeft, PrefixInput) 32 | 33 | switch prefix { 34 | case PrefixBlockStart, PrefixBlockEnd: 35 | padLeft = "" 36 | end = fmt.Sprintf(" \n") 37 | case PrefixInput: 38 | end = "" 39 | case PrefixRaw: 40 | fmt.Printf("\r%s\n", strings.Repeat(" ", 16)) 41 | for _, l := range strings.Split(fmt.Sprintf(format, c...), "\n") { 42 | fmt.Printf("%s%s\n", strings.Repeat(" ", 4), l) 43 | } 44 | fmt.Print("\n") 45 | 46 | return 47 | } 48 | 49 | fmt.Printf("\r%s%s %s%s", padLeft, prefix, fmt.Sprintf(format, c...), end) 50 | } 51 | -------------------------------------------------------------------------------- /scripts/server.lua: -------------------------------------------------------------------------------- 1 | local stat = require("stat") 2 | local runtime = require("runtime") 3 | local ws = require("ws") 4 | local time = require("time") 5 | 6 | if runtime.isMaster() then 7 | local start = time.now(time.ms) 8 | 9 | stat.new("threads", stat.abs()) 10 | stat.new("duration", stat.abs({ measure = "ms" })) 11 | stat.new("connections", stat.abs()) 12 | stat.new("messages", stat.abs()) 13 | 14 | for i = 0, (runtime.numCPU-1) do 15 | runtime.fork() 16 | stat.add("threads", 1) 17 | end 18 | 19 | runtime.on("exit", function() 20 | print("bye!") 21 | stat.add("duration", time.now(time.ms) - start) 22 | end) 23 | else 24 | local server = ws.createServer() 25 | local addr = runtime.get("listen") 26 | 27 | server.listen(addr, function(conn) 28 | stat.add("connections", 1) 29 | conn.listen(function() 30 | stat.add("messages", 1) 31 | end) 32 | end) 33 | 34 | server.on("error", function(err) 35 | print("error: ", err) 36 | end) 37 | 38 | server.on("listening", function() 39 | print("listening: ", addr) 40 | end) 41 | end -------------------------------------------------------------------------------- /lua/util/util.go: -------------------------------------------------------------------------------- 1 | package util 2 | 3 | import ( 4 | "github.com/yuin/gopher-lua" 5 | "net/http" 6 | ) 7 | 8 | func HeadersToMap(h http.Header) map[string]string { 9 | m := make(map[string]string) 10 | for key := range h { 11 | m[key] = h.Get(key) 12 | } 13 | return m 14 | } 15 | 16 | func MapOfStringToTable(L *lua.LState, m map[string]string) *lua.LTable { 17 | table := L.NewTable() 18 | for key, value := range m { 19 | table.RawSetString(key, lua.LString(value)) 20 | } 21 | return table 22 | } 23 | 24 | func MapOfInterfaceToTable(L *lua.LState, m map[string]interface{}) *lua.LTable { 25 | table := L.NewTable() 26 | for key, value := range m { 27 | switch v := value.(type) { 28 | case string: 29 | table.RawSetString(key, lua.LString(v)) 30 | case int: 31 | table.RawSetString(key, lua.LNumber(v)) 32 | case uint: 33 | table.RawSetString(key, lua.LNumber(v)) 34 | case float64: 35 | table.RawSetString(key, lua.LNumber(v)) 36 | } 37 | } 38 | return table 39 | } 40 | 41 | func HeadersFromTable(t *lua.LTable) (h http.Header) { 42 | h = make(http.Header, t.Len()) 43 | t.ForEach(func(k, v lua.LValue) { 44 | h.Set(k.String(), v.String()) 45 | }) 46 | return 47 | } 48 | -------------------------------------------------------------------------------- /cli/input/input.go: -------------------------------------------------------------------------------- 1 | package input 2 | 3 | import ( 4 | "github.com/chzyer/readline" 5 | "io" 6 | ) 7 | 8 | type Message struct { 9 | Err error 10 | Data []byte 11 | } 12 | 13 | func ReadLine(cfg *readline.Config) (r []byte, err error) { 14 | rl, err := readline.NewEx(cfg) 15 | if err != nil { 16 | return 17 | } 18 | defer rl.Close() 19 | 20 | line, err := rl.Readline() 21 | if err != nil { 22 | return 23 | } 24 | 25 | return []byte(line), nil 26 | } 27 | 28 | func ReadLineAsync(done <-chan struct{}, cfg *readline.Config) (<-chan Message, error) { 29 | ch := make(chan Message) 30 | 31 | rl, err := readline.NewEx(cfg) 32 | if err != nil { 33 | return nil, err 34 | } 35 | 36 | go func() { 37 | for { 38 | var msg Message 39 | line, err := rl.Readline() 40 | if err != nil { 41 | if err == readline.ErrInterrupt { 42 | msg = Message{Err: io.EOF} 43 | } else { 44 | msg = Message{Err: err} 45 | } 46 | 47 | } else { 48 | msg = Message{Data: []byte(line)} 49 | } 50 | 51 | select { 52 | case <-done: 53 | // 54 | default: 55 | ch <- msg 56 | if msg.Err != nil { 57 | rl.Close() 58 | return 59 | } 60 | } 61 | } 62 | }() 63 | 64 | return ch, nil 65 | } 66 | -------------------------------------------------------------------------------- /gws.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "flag" 5 | "fmt" 6 | "github.com/gobwas/gws/cli/color" 7 | "github.com/gobwas/gws/client" 8 | "github.com/gobwas/gws/config" 9 | "github.com/gobwas/gws/lua" 10 | "github.com/gobwas/gws/server" 11 | "io" 12 | "os" 13 | "strings" 14 | ) 15 | 16 | const ( 17 | modeServer = "server" 18 | modeClient = "client" 19 | modeScript = "script" 20 | ) 21 | 22 | var modes = []string{modeServer, modeClient, modeScript} 23 | 24 | func main() { 25 | flag.Usage = func() { 26 | fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0]) 27 | fmt.Fprintf(os.Stderr, "%s %s|%s|%s [options]\n", os.Args[0], modeClient, modeServer, modeScript) 28 | fmt.Fprintf(os.Stderr, "options:\n") 29 | flag.PrintDefaults() 30 | } 31 | if len(os.Args) < 2 { 32 | flag.Usage() 33 | os.Exit(1) 34 | } 35 | flag.CommandLine.Parse(os.Args[2:]) 36 | 37 | cfg, err := config.Parse() 38 | if err != nil { 39 | flag.Usage() 40 | os.Exit(1) 41 | } 42 | 43 | switch os.Args[1] { 44 | case modeServer: 45 | err = server.Go(cfg) 46 | case modeClient: 47 | err = client.Go(cfg) 48 | case modeScript: 49 | err = lua.Go(cfg) 50 | default: 51 | err = fmt.Errorf("mode is required to be a one of `%s`; but `%s` given", color.Cyan(strings.Join(modes, "`, `")), color.Yellow(os.Args[1])) 52 | } 53 | 54 | if err != nil && err != io.EOF { 55 | fmt.Fprintf(os.Stderr, fmt.Sprintf("%s %s\n\n", color.Red("error:"), err)) 56 | os.Exit(1) 57 | } 58 | 59 | fmt.Printf("\r") 60 | os.Exit(0) 61 | } 62 | -------------------------------------------------------------------------------- /display/display_test.go: -------------------------------------------------------------------------------- 1 | package display_test 2 | 3 | import ( 4 | "fmt" 5 | "github.com/gobwas/gws/display" 6 | "math/rand" 7 | "os" 8 | "strings" 9 | "testing" 10 | "time" 11 | ) 12 | 13 | func TestDisplay(t *testing.T) { 14 | d := display.NewDisplay(os.Stderr, display.Config{ 15 | TabSize: 4, 16 | Interval: time.Millisecond * 500, 17 | }) 18 | r1 := d.Row() 19 | r1.Column(display.Column{ 20 | Width: 20, 21 | Height: 2, 22 | Content: func() string { 23 | return fmt.Sprintf("left col: %s", strings.Repeat("*", rand.Intn(10))) 24 | }, 25 | }) 26 | r1.Column(display.Column{ 27 | Width: 20, 28 | Height: 2, 29 | Content: func() string { 30 | return fmt.Sprintf("right col: %s", strings.Repeat("*", rand.Intn(10))) 31 | }, 32 | }) 33 | 34 | r2 := d.Row() 35 | r2.Column(display.Column{ 36 | Height: 2, 37 | Width: 100, 38 | Content: func() string { 39 | return fmt.Sprintf("long middle single col: %s", strings.Repeat("*", rand.Intn(10))) 40 | }, 41 | }) 42 | 43 | r3 := d.Row() 44 | var multilineA []string 45 | r3.Column(display.Column{ 46 | Width: 20, 47 | Height: 4, 48 | Content: func() string { 49 | multilineA = append(multilineA, fmt.Sprintf("left counter: %d", rand.Intn(100))) 50 | return strings.Join(multilineA, "\n") 51 | }, 52 | }) 53 | var multilineB []string 54 | r3.Column(display.Column{ 55 | Width: 20, 56 | Height: 10, 57 | Content: func() string { 58 | multilineB = append(multilineB, fmt.Sprintf("right counter: %d", rand.Intn(100))) 59 | return strings.Join(multilineB, "\n") 60 | }, 61 | }) 62 | d.On() 63 | time.Sleep(time.Minute) 64 | } 65 | -------------------------------------------------------------------------------- /ws/connection.go: -------------------------------------------------------------------------------- 1 | package ws 2 | 3 | import ( 4 | "errors" 5 | "github.com/gorilla/websocket" 6 | "sync" 7 | ) 8 | 9 | type Connection struct { 10 | once sync.Once 11 | 12 | conn *websocket.Conn 13 | done chan struct{} 14 | in chan ReceiveRequest 15 | out chan WriteRequest 16 | running bool 17 | } 18 | 19 | func NewConnection(c *websocket.Conn) *Connection { 20 | return &Connection{ 21 | conn: c, 22 | done: make(chan struct{}), 23 | in: make(chan ReceiveRequest), 24 | out: make(chan WriteRequest), 25 | } 26 | } 27 | 28 | func (c *Connection) SendAsync(msg MessageRaw) chan error { 29 | resp := make(chan error) 30 | c.out <- WriteRequest{msg, resp} 31 | return resp 32 | } 33 | 34 | func (c *Connection) Send(msg MessageRaw) error { 35 | return <-c.SendAsync(msg) 36 | } 37 | 38 | func (c *Connection) ReceiveAsync() chan MessageAndError { 39 | resp := make(chan MessageAndError) 40 | c.in <- ReceiveRequest{resp} 41 | return resp 42 | } 43 | 44 | func (c *Connection) Receive() (MessageRaw, error) { 45 | result := <-c.ReceiveAsync() 46 | return result.Message, result.Error 47 | } 48 | 49 | func (c *Connection) Done() <-chan struct{} { 50 | return c.done 51 | } 52 | 53 | func (c *Connection) IsRunning() bool { 54 | return c.running 55 | } 56 | 57 | func (c *Connection) InitIOWorkers() { 58 | c.once.Do(func() { 59 | WriteToConnFromChan(c.done, c.conn, c.out) 60 | ReadFromConnToChan(c.done, c.conn, c.in) 61 | c.running = true 62 | }) 63 | } 64 | 65 | func (c *Connection) Close() error { 66 | select { 67 | case <-c.done: 68 | return errors.New("already closed") 69 | default: 70 | close(c.done) 71 | close(c.in) 72 | close(c.out) 73 | } 74 | 75 | return c.conn.Close() 76 | } 77 | -------------------------------------------------------------------------------- /lua/mod/ws/server.go: -------------------------------------------------------------------------------- 1 | package ws 2 | 3 | import ( 4 | "github.com/gobwas/gws/ev" 5 | "github.com/gobwas/gws/lua/mod" 6 | "github.com/gobwas/gws/ws" 7 | "github.com/yuin/gopher-lua" 8 | ) 9 | 10 | type Server struct { 11 | loop *ev.Loop 12 | emitter *mod.Emitter 13 | config ws.ServerConfig 14 | } 15 | 16 | func NewServer(l *ev.Loop, c ws.ServerConfig) *Server { 17 | return &Server{ 18 | loop: l, 19 | emitter: mod.NewEmitter(), 20 | config: c, 21 | } 22 | } 23 | 24 | func (s *Server) Emit(name string, args ...interface{}) { 25 | s.loop.Call(func() { 26 | s.emitter.Emit(name, args...) 27 | }) 28 | } 29 | 30 | func (s *Server) ToTable(L *lua.LState) *lua.LTable { 31 | table := L.NewTable() 32 | 33 | table.RawSetString("listen", L.NewClosure(func(L *lua.LState) int { 34 | addr := L.ToString(1) 35 | cb := L.ToFunction(2) 36 | 37 | req := s.config 38 | req.Addr = addr 39 | 40 | var e error // hack to catch sync error 41 | s.loop.Request(101, req, func(err error, msg interface{}) { 42 | if err != nil { 43 | e = err 44 | s.Emit("error", err.Error()) 45 | return 46 | } 47 | 48 | if c, ok := msg.(*ws.Connection); ok { 49 | c.InitIOWorkers() 50 | conn := NewConn(c, s.loop) 51 | L.CallByParam(lua.P{ 52 | Fn: cb, 53 | NRet: 0, 54 | Protect: false, 55 | }, conn.ToTable(L)) 56 | return 57 | } 58 | 59 | panic("ws: unexpected data in request callback") 60 | }) 61 | 62 | if e != nil { // sync error 63 | L.Push(lua.LString(e.Error())) 64 | return 1 65 | } else { 66 | s.Emit("listening") 67 | } 68 | 69 | return 0 70 | })) 71 | 72 | table.RawSetString("on", s.emitter.ExportOn(L)) 73 | table.RawSetString("off", s.emitter.ExportOff(L)) 74 | 75 | return table 76 | } 77 | -------------------------------------------------------------------------------- /scripts/load.lua: -------------------------------------------------------------------------------- 1 | local stat = require("stat") 2 | local time = require("time") 3 | local runtime = require("runtime") 4 | local ws = require("ws") 5 | 6 | local start = time.now(time.ms) 7 | 8 | if runtime.isMaster() then 9 | stat.new("err", stat.abs()) 10 | stat.new("threads", stat.abs()) 11 | stat.new("sent", stat.per("1s"), stat.abs()) 12 | stat.new("delay", stat.avg({ measure = "ms" })) 13 | stat.new("duration", stat.abs({ measure = "ms" })) 14 | 15 | for i = 0, (runtime.numCPU-1) do 16 | runtime.fork() 17 | stat.add("threads", 1) 18 | end 19 | 20 | runtime.on("exit", function() 21 | print("bye!") 22 | stat.add("duration", time.now(time.ms) - start) 23 | end) 24 | else 25 | --[[ 26 | runtime exports: 27 | url -u flag value 28 | headers -h flag value 29 | id incremental identifier of the fork 30 | ]] 31 | ws.connect({ url = runtime.get("url") }, function(err, conn) 32 | if (err ~= nil) then 33 | print("could not connect: ", err) 34 | return 35 | else 36 | -- if you want to get the maximum throughput in your tests 37 | -- and not need to receive answers from the server under test, 38 | -- it is preferred to be synchronous. 39 | while true do 40 | local start = time.now(time.ms) 41 | if (conn.send("hello") ~= nil) then 42 | stat.add("err", 1) 43 | else 44 | stat.add("sent", 1) 45 | stat.add("delay", time.now(time.ms) - start) 46 | end 47 | end 48 | end 49 | end) 50 | 51 | runtime.on("exit", function() 52 | print("bye!") 53 | end) 54 | end 55 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # :sparkles: gws 2 | 3 | > CLI tool for websocket testing 4 | 5 | ![Demo](https://cdn.rawgit.com/gobwas/gws/static/demo.gif) 6 | 7 | ## Install 8 | 9 | ```shell 10 | go get github.com/gobwas/gws 11 | ``` 12 | 13 | ## Usage exmaples 14 | 15 | Connect to the websocket server: 16 | 17 | ```shell 18 | gws client -url="ws://my.cool.address" 19 | ``` 20 | 21 | Run simple server and type response messages in terminal: 22 | 23 | ```shell 24 | gws server -listen=":8888" -response=prompt 25 | ``` 26 | 27 | Or just simple echo: 28 | 29 | ```shell 30 | gws server -listen=":8888" -response=echo 31 | ``` 32 | 33 | Run lua script: 34 | 35 | ```shell 36 | gws script -path=./my_cool_script.lua 37 | ``` 38 | 39 | Usage info: 40 | 41 | ```shell 42 | Usage of gws: 43 | gws client|server|script [options] 44 | options: 45 | -header string 46 | list of headers to be passed during handshake (both in client or server) 47 | format: 48 | { pair[ ";" pair...] }, 49 | pair: 50 | { key ":" value } 51 | -listen string 52 | address to listen (default ":3000") 53 | -origin string 54 | use this glob pattern for server origin checks 55 | -path string 56 | path to lua script 57 | -response value 58 | how should server response on message (echo, mirror, prompt, null) (default null) 59 | -retry int 60 | try to reconnect x times (default 1) 61 | -statd duration 62 | server statistics dump interval (default 1s) 63 | -url string 64 | address to connect (default ":3000") 65 | -verbose 66 | verbose output 67 | ``` 68 | 69 | ## Scripting 70 | 71 | gws brings you ability to implement your tests logic in `.lua` scripts. 72 | Please look at `scripts` folder in this repository to find an examples of scripting. 73 | 74 | ## Why 75 | 76 | `gws` is highly inspired by [wsd](https://github.com/alexanderGugel/wsd) and [iocat](https://github.com/moul/iocat). But in both 77 | tools I found not existing features that I was needed some how. 78 | -------------------------------------------------------------------------------- /ev/timer/timer.go: -------------------------------------------------------------------------------- 1 | package timer 2 | 3 | import ( 4 | "errors" 5 | "github.com/gobwas/gws/ev" 6 | "sync/atomic" 7 | "time" 8 | ) 9 | 10 | type delay struct { 11 | timer *time.Timer 12 | timeout *Timeout 13 | loop *ev.Loop 14 | cb ev.Callback 15 | } 16 | 17 | type Handler struct { 18 | t ev.RequestType 19 | delay chan delay 20 | count int32 21 | } 22 | 23 | func NewHandler(t ev.RequestType) *Handler { 24 | return &Handler{ 25 | t: t, 26 | delay: make(chan delay, 1), 27 | } 28 | } 29 | 30 | func (h *Handler) Init() { 31 | go func() { 32 | for { 33 | for d := range h.delay { 34 | if d.timeout.dropped { 35 | d.timer.Stop() 36 | atomic.AddInt32(&h.count, -1) 37 | continue 38 | } 39 | 40 | select { 41 | case t := <-d.timer.C: 42 | d.loop.Call(func() { d.cb(nil, t) }) 43 | if d.timeout.repeat { 44 | d.timer.Reset(d.timeout.delay) 45 | h.delay <- d 46 | } else { 47 | atomic.AddInt32(&h.count, -1) 48 | } 49 | default: 50 | h.delay <- d 51 | } 52 | } 53 | } 54 | }() 55 | } 56 | 57 | func (h *Handler) SetTimeout(loop *ev.Loop, t *Timeout, cb ev.Callback) { 58 | h.delay <- delay{ 59 | cb: cb, 60 | loop: loop, 61 | timer: time.NewTimer(t.delay), 62 | timeout: t, 63 | } 64 | atomic.AddInt32(&h.count, 1) 65 | } 66 | 67 | func (h *Handler) Handle(loop *ev.Loop, data interface{}, cb ev.Callback) error { 68 | t, ok := data.(*Timeout) 69 | if !ok { 70 | return errors.New("unexpected data") 71 | } 72 | h.SetTimeout(loop, t, cb) 73 | 74 | return nil 75 | } 76 | 77 | func (h *Handler) IsActive(loop *ev.Loop) bool { 78 | return atomic.LoadInt32(&h.count) > 0 79 | } 80 | 81 | type Timeout struct { 82 | delay time.Duration 83 | dropped bool 84 | repeat bool 85 | } 86 | 87 | func NewTimeout(t time.Duration) *Timeout { 88 | return &Timeout{ 89 | delay: t, 90 | } 91 | } 92 | 93 | func NewTicker(t time.Duration) *Timeout { 94 | return &Timeout{ 95 | delay: t, 96 | repeat: true, 97 | } 98 | } 99 | 100 | func (t *Timeout) Stop() { 101 | t.dropped = true 102 | } 103 | -------------------------------------------------------------------------------- /lua/mod/runtime/runtime.go: -------------------------------------------------------------------------------- 1 | package runtime 2 | 3 | import ( 4 | "github.com/gobwas/gws/ev" 5 | "github.com/gobwas/gws/lua/mod" 6 | "github.com/yuin/gopher-lua" 7 | "runtime" 8 | "sync/atomic" 9 | "time" 10 | ) 11 | 12 | type Runtime struct { 13 | emitter *mod.Emitter 14 | storage *mod.Storage 15 | 16 | exported int32 17 | initTime time.Time 18 | loop *ev.Loop 19 | fork forkFn 20 | } 21 | 22 | type callback struct { 23 | event string 24 | fn *lua.LFunction 25 | } 26 | 27 | type forkFn func() error 28 | 29 | func New(loop *ev.Loop) *Runtime { 30 | return &Runtime{ 31 | emitter: mod.NewEmitter(), 32 | storage: mod.NewStorage(), 33 | initTime: time.Now(), 34 | loop: loop, 35 | } 36 | } 37 | 38 | func (m *Runtime) SetForkFn(f forkFn) { 39 | if atomic.LoadInt32(&m.exported) > 0 { 40 | panic("could not set fork function after runtime is exported") 41 | } 42 | m.fork = f 43 | } 44 | 45 | func (m *Runtime) Emit(name string) { 46 | m.loop.Call(func() { 47 | m.emitter.Emit(name) 48 | }) 49 | } 50 | 51 | func (m *Runtime) Set(key string, value interface{}) { 52 | m.storage.Set(key, value) 53 | } 54 | 55 | func (m *Runtime) Get(key string) (result interface{}) { 56 | return m.storage.Get(key) 57 | } 58 | 59 | func (m *Runtime) Exports() lua.LGFunction { 60 | if atomic.LoadInt32(&m.exported) > 0 { 61 | panic("runtime could be exported only once") 62 | } else { 63 | atomic.AddInt32(&m.exported, 1) 64 | } 65 | 66 | return func(L *lua.LState) int { 67 | mod := L.NewTable() 68 | 69 | if m.fork != nil { 70 | mod.RawSetString("fork", L.NewClosure(func(L *lua.LState) int { 71 | if err := m.fork(); err != nil { 72 | L.Push(lua.LString(err.Error())) 73 | return 1 74 | } 75 | return 0 76 | })) 77 | } 78 | 79 | mod.RawSetString("isMaster", L.NewClosure(func(L *lua.LState) int { 80 | L.Push(lua.LBool(m.fork != nil)) 81 | return 1 82 | })) 83 | 84 | mod.RawSetString("numCPU", lua.LNumber(runtime.NumCPU())) 85 | 86 | mod.RawSetString("set", m.storage.ExportSet(L)) 87 | mod.RawSetString("get", m.storage.ExportGet(L)) 88 | 89 | mod.RawSetString("on", m.emitter.ExportOn(L)) 90 | mod.RawSetString("off", m.emitter.ExportOff(L)) 91 | 92 | L.Push(mod) 93 | return 1 94 | } 95 | } 96 | -------------------------------------------------------------------------------- /client/client.go: -------------------------------------------------------------------------------- 1 | package client 2 | 3 | import ( 4 | "flag" 5 | "github.com/chzyer/readline" 6 | "github.com/gobwas/gws/cli" 7 | "github.com/gobwas/gws/cli/color" 8 | cliInput "github.com/gobwas/gws/cli/input" 9 | "github.com/gobwas/gws/config" 10 | "github.com/gobwas/gws/util" 11 | "github.com/gobwas/gws/ws" 12 | "github.com/gorilla/websocket" 13 | "io" 14 | "net/http" 15 | "time" 16 | ) 17 | 18 | var limit = flag.Int("retry", 1, "try to reconnect x times") 19 | 20 | const readLineTemp = "/tmp/gws_readline_client.tmp" 21 | 22 | func Go(c config.Config) error { 23 | var conn *websocket.Conn 24 | var err error 25 | 26 | for i := 0; i < *limit; i++ { 27 | conn, err = getConn(c.URI, c.Headers) 28 | if err == nil { 29 | break 30 | } 31 | time.Sleep(time.Millisecond * time.Duration(100*i)) 32 | } 33 | if err != nil { 34 | cli.Printf(cli.PrefixTheEnd, "could not connect: %s", color.Red(err)) 35 | return err 36 | } 37 | 38 | done := make(chan struct{}) 39 | output, err := cliInput.ReadLineAsync(done, &readline.Config{ 40 | Prompt: cli.PaddingLeft + "> ", 41 | HistoryFile: readLineTemp, 42 | }) 43 | if err != nil { 44 | return err 45 | } 46 | input := ws.ReadAsyncFromConn(done, conn) 47 | 48 | for { 49 | select { 50 | case in := <-input: 51 | if in.Err != nil { 52 | if in.Err == io.EOF { 53 | cli.Printf(cli.PrefixTheEnd, "%s %s", color.Magenta(in.Err), color.Red("server has closed connection")) 54 | } else { 55 | cli.Printf(cli.PrefixInfo, "%s %s", color.Magenta(in.Err), color.Red("unknown error")) 56 | } 57 | 58 | cli.Printf(cli.PrefixBlockEnd, "") 59 | return in.Err 60 | } 61 | 62 | cli.Printf(cli.PrefixIncoming, "%s: %s", color.Magenta(in.Kind), color.Cyan(string(in.Data))) 63 | 64 | case out := <-output: 65 | if out.Err != nil { 66 | cli.Printf(cli.PrefixTheEnd, "%s %s", color.Magenta(out.Err), color.Red("input closed")) 67 | return out.Err 68 | } 69 | 70 | err := ws.WriteToConn(conn, ws.TextMessage, out.Data) 71 | if err != nil { 72 | cli.Printf(cli.PrefixInfo, "%s", color.Red(err)) 73 | } 74 | } 75 | } 76 | } 77 | 78 | func getConn(uri string, h http.Header) (*websocket.Conn, error) { 79 | conn, resp, err := ws.GetConn(uri, h) 80 | if config.Verbose { 81 | req, res, _ := util.DumpRequestResponse(resp) 82 | cli.Printf(cli.PrefixRaw, "%s", color.Green(string(req))) 83 | cli.Printf(cli.PrefixRaw, "%s", color.Cyan(string(res))) 84 | } 85 | if err != nil { 86 | cli.Printf(cli.PrefixInfo, "%s %s", color.Magenta(err), color.Red("could not connect")) 87 | return nil, err 88 | } 89 | 90 | cli.Printf(cli.PrefixInfo, "connected to %s", color.Green(uri)) 91 | cli.Printf(cli.PrefixEmpty, "") 92 | 93 | return conn, nil 94 | } 95 | -------------------------------------------------------------------------------- /scripts/client.lua: -------------------------------------------------------------------------------- 1 | local stat = require("stat") 2 | local time = require("time") 3 | local runtime = require("runtime") 4 | local ws = require("ws") 5 | 6 | local start = time.now(time.ms) 7 | 8 | if runtime.isMaster() then 9 | stat.new("send_err", stat.abs()) 10 | stat.new("send_ok", stat.abs()) 11 | stat.new("recv_err", stat.abs()) 12 | stat.new("recv_ok", stat.abs()) 13 | stat.new("threads", stat.abs()) 14 | stat.new("recv_delay", stat.avg({ measure = "ms" })) 15 | stat.new("duration", stat.abs({ measure = "ms" })) 16 | 17 | for i = 0, (runtime.numCPU-1) do 18 | runtime.fork() 19 | stat.add("threads", 1) 20 | end 21 | 22 | runtime.on("exit", function() 23 | print("bye!") 24 | stat.add("duration", time.now(time.ms) - start) 25 | end) 26 | else 27 | local function send(conn) 28 | conn.send("hello", function(err) 29 | if (err ~= nil) then 30 | stat.add("err", 1) 31 | else 32 | stat.add("sent", 1) 33 | stat.add("delay", time.now(time.ms) - start) 34 | end 35 | end) 36 | end 37 | 38 | --[[ 39 | runtime exports: 40 | url -u flag value 41 | headers -h flag value 42 | id incremental identifier of the fork 43 | ]] 44 | ws.connect({ url = runtime.get("url") }, function(err, conn) 45 | if (err ~= nil) then 46 | print("could not connect: ", err) 47 | else 48 | print("connect ok") 49 | 50 | local start; 51 | local send = function() 52 | start = time.now(time.ms) 53 | conn.send("hello", function(err) 54 | if (err ~= nil) then 55 | stat.add("send_err", 1) 56 | print("send error: ", err) 57 | else 58 | stat.add("send_ok", 1) 59 | print("send ok") 60 | end 61 | end) 62 | end 63 | 64 | conn.listen(function(err, msg) 65 | if err ~= nil then 66 | stat.add("recv_err", 1) 67 | print("receive error: ", err) 68 | else 69 | stat.add("recv_ok", 1) 70 | stat.add("recv_delay", time.now(time.ms) - start) 71 | print(string.format("receive ok: %s (%d)", msg)) 72 | end 73 | 74 | time.setTimeout(100, send) 75 | end) 76 | 77 | conn.on("close", function() 78 | print("conn closed") 79 | end) 80 | 81 | send() 82 | end 83 | end) 84 | 85 | runtime.on("exit", function() 86 | print("bye!") 87 | end) 88 | end 89 | -------------------------------------------------------------------------------- /ev/ws/server.go: -------------------------------------------------------------------------------- 1 | package ws 2 | 3 | import ( 4 | "fmt" 5 | "github.com/gobwas/gws/ev" 6 | "github.com/gobwas/gws/ws" 7 | "github.com/gorilla/websocket" 8 | "reflect" 9 | "sync" 10 | "sync/atomic" 11 | ) 12 | 13 | type serverDesc struct { 14 | cfg ws.ServerConfig 15 | server *ws.Server 16 | loops []*ev.Loop 17 | } 18 | 19 | type ServerHandler struct { 20 | mu sync.Mutex 21 | pending map[*ev.Loop]int32 22 | servers map[string]serverDesc 23 | stop chan struct{} 24 | stopped int32 25 | } 26 | 27 | func NewServerHandler() *ServerHandler { 28 | return &ServerHandler{ 29 | pending: make(map[*ev.Loop]int32), 30 | servers: make(map[string]serverDesc), 31 | stop: make(chan struct{}), 32 | } 33 | } 34 | 35 | func (h *ServerHandler) Init(*ev.Loop) error { 36 | return nil 37 | } 38 | 39 | func (h *ServerHandler) Handle(loop *ev.Loop, data interface{}, cb ev.Callback) error { 40 | switch v := data.(type) { 41 | case ws.ServerConfig: 42 | h.doListen(loop, v, cb) 43 | default: 44 | return fmt.Errorf("unknown request format to ws handler: %s", data) 45 | } 46 | return nil 47 | } 48 | 49 | func (h *ServerHandler) Stop() { 50 | if atomic.CompareAndSwapInt32(&h.stopped, 0, 1) { 51 | close(h.stop) 52 | } 53 | } 54 | 55 | func (h *ServerHandler) IsActive(loop *ev.Loop) bool { 56 | return h.getPending(loop) > 0 57 | } 58 | 59 | func (h *ServerHandler) doListen(loop *ev.Loop, cfg ws.ServerConfig, cb ev.Callback) { 60 | h.mu.Lock() 61 | defer h.mu.Unlock() 62 | 63 | desc, ok := h.servers[cfg.Addr] 64 | if ok { 65 | if !reflect.DeepEqual(desc.cfg, cfg) { 66 | cb(fmt.Errorf("already listening on %s with different configuration", cfg.Addr), nil) 67 | return 68 | } 69 | for _, l := range desc.loops { 70 | if l == loop { 71 | cb(fmt.Errorf("already listening on %s in current thread", cfg.Addr), nil) 72 | return 73 | } 74 | } 75 | } else { 76 | s := ws.NewServer(ws.ServerConfig{ 77 | Key: cfg.Key, 78 | Cert: cfg.Cert, 79 | Addr: cfg.Addr, 80 | Headers: cfg.Headers, 81 | Origin: cfg.Origin, 82 | }) 83 | defer s.Listen(h.stop) 84 | desc = serverDesc{server: s, cfg: cfg} 85 | h.servers[cfg.Addr] = desc 86 | } 87 | 88 | desc.loops = append(desc.loops, loop) 89 | desc.server.Handle(ws.HandlerFunc(func(conn *websocket.Conn, err error) { 90 | if err != nil { 91 | loop.Call(func() { cb(err, nil) }) 92 | } else { 93 | loop.Call(func() { cb(nil, ws.NewConnection(conn)) }) 94 | } 95 | })) 96 | 97 | h.pending[loop] += 1 98 | desc.server.Defer(func() { 99 | h.mu.Lock() 100 | defer h.mu.Unlock() 101 | h.pending[loop] -= 1 102 | }) 103 | } 104 | 105 | func (h *ServerHandler) getPending(loop *ev.Loop) int32 { 106 | h.mu.Lock() 107 | defer h.mu.Unlock() 108 | return h.pending[loop] 109 | } 110 | -------------------------------------------------------------------------------- /bufio/writer.go: -------------------------------------------------------------------------------- 1 | // Package bufio brings tools for io. 2 | // It extends standard bufio package with prefixed writer and ring buffer. 3 | package bufio 4 | 5 | import "io" 6 | 7 | // PrefixWriter adds prefix to every write call. 8 | type PrefixWriter struct { 9 | dest io.Writer 10 | prefix string 11 | } 12 | 13 | func NewPrefixWriter(dest io.Writer, prefix string) *PrefixWriter { 14 | return &PrefixWriter{dest, prefix} 15 | } 16 | 17 | // Write writes p into underlying writer with prefix. 18 | func (w PrefixWriter) Write(p []byte) (int, error) { 19 | ret := make([]byte, len(p)+len(w.prefix)) 20 | ret = append(ret, w.prefix...) 21 | ret = append(ret, p...) 22 | return w.dest.Write(ret) 23 | } 24 | 25 | // RingBufferWriter implements ring buffer for io operations. 26 | type RingBufferWriter struct { 27 | ring *ring 28 | dest io.Writer 29 | } 30 | 31 | func NewWriter(dest io.Writer, size int) *RingBufferWriter { 32 | return &RingBufferWriter{ 33 | dest: dest, 34 | ring: newRing(size), 35 | } 36 | } 37 | 38 | // Write writes p to the underlying circular buffer. 39 | // It returns len(p) and optional error. 40 | func (w *RingBufferWriter) Write(p []byte) (int, error) { 41 | w.ring.append(p...) 42 | return len(p), nil 43 | } 44 | 45 | // Flush dumps contents of circular buffer into the underlying io.Writer. 46 | // Flush is the same as Dump, except the thing, that flush drops buffer contents. 47 | func (w *RingBufferWriter) Flush() (err error) { 48 | _, err = w.dest.Write(w.ring.flush()) 49 | return 50 | } 51 | 52 | // Dump dumps contents of circular buffer into the underlying io.Writer. 53 | // Dump do not drops buffer contents. 54 | func (w *RingBufferWriter) Dump() (err error) { 55 | _, err = w.dest.Write(w.ring.dump()) 56 | return 57 | } 58 | 59 | type ring struct { 60 | data []byte 61 | size int 62 | pos int 63 | len int 64 | } 65 | 66 | func newRing(size int) *ring { 67 | return &ring{ 68 | size: size, 69 | data: make([]byte, size), 70 | } 71 | } 72 | 73 | func (r *ring) append(b ...byte) { 74 | var start int 75 | if len(b) > r.size { 76 | // get bytes that could be stored 77 | // from the end of b 78 | start = len(b) - r.size 79 | } 80 | 81 | var end int 82 | if len(b) < r.size-r.pos { 83 | end = len(b) 84 | } else { 85 | end = r.size - r.pos 86 | } 87 | 88 | rc := copy(r.data[r.pos:], b[start:end]) 89 | lc := copy(r.data[:r.pos], b[start+end:]) 90 | 91 | r.pos = (r.pos + rc + lc) % r.size 92 | if r.len+rc+lc > r.size { 93 | r.len = r.size 94 | } else { 95 | r.len += rc + lc 96 | } 97 | } 98 | 99 | func (r *ring) flush() (ret []byte) { 100 | ret = r.dump() 101 | r.len = 0 102 | r.pos = 0 103 | return 104 | } 105 | 106 | func (r *ring) dump() (ret []byte) { 107 | ret = make([]byte, r.len) 108 | var start int 109 | if r.len == r.size { 110 | copy(ret, r.data[r.pos:]) 111 | start = r.size - r.pos 112 | } 113 | copy(ret[start:], r.data[:r.pos]) 114 | return 115 | } 116 | -------------------------------------------------------------------------------- /lua/mod/ws/ws.go: -------------------------------------------------------------------------------- 1 | package ws 2 | 3 | import ( 4 | "github.com/gobwas/gws/ev" 5 | evws "github.com/gobwas/gws/ev/ws" 6 | luautil "github.com/gobwas/gws/lua/util" 7 | "github.com/gobwas/gws/ws" 8 | "github.com/yuin/gopher-lua" 9 | "net/http" 10 | ) 11 | 12 | type Mod struct { 13 | loop *ev.Loop 14 | } 15 | 16 | func New(loop *ev.Loop) *Mod { 17 | return &Mod{ 18 | loop: loop, 19 | } 20 | } 21 | 22 | func (m *Mod) Exports() lua.LGFunction { 23 | return func(L *lua.LState) int { 24 | mod := L.NewTable() 25 | 26 | mod.RawSetString("createServer", L.NewClosure(func(L *lua.LState) int { 27 | var cfg ws.ServerConfig 28 | if opts := L.ToTable(1); opts != nil { 29 | opts.ForEach(func(key lua.LValue, value lua.LValue) { 30 | if key.Type() == lua.LTString { 31 | switch key.String() { 32 | case "cert": 33 | cfg.Cert = value.String() 34 | case "key": 35 | cfg.Key = value.String() 36 | case "origin": 37 | cfg.Origin = value.String() 38 | case "headers": 39 | t, ok := value.(*lua.LTable) 40 | if ok { 41 | cfg.Headers = luautil.HeadersFromTable(t) 42 | } 43 | } 44 | } 45 | }) 46 | } 47 | 48 | server := NewServer(m.loop, cfg) 49 | L.Push(server.ToTable(L)) 50 | 51 | return 1 52 | })) 53 | 54 | mod.RawSetString("connect", L.NewClosure(func(L *lua.LState) int { 55 | opts := L.ToTable(1) 56 | cb := L.ToFunction(2) 57 | 58 | var uri string 59 | if u := opts.RawGetString("url"); u.Type() == lua.LTString { 60 | uri = u.String() 61 | } else { 62 | m.loop.Call(func() { 63 | L.CallByParam(lua.P{ 64 | Fn: cb, 65 | NRet: 0, 66 | Protect: false, 67 | }, lua.LString("url is expected to be a string in options table"), lua.LNil) 68 | }) 69 | return 0 70 | } 71 | 72 | var headers http.Header 73 | if h := opts.RawGetString("headers"); h.Type() == lua.LTTable { 74 | headers = make(http.Header) 75 | h.(*lua.LTable).ForEach(func(k, v lua.LValue) { 76 | if k.Type() == lua.LTString && v.Type() == lua.LTString { 77 | headers.Set(k.String(), v.String()) 78 | } 79 | }) 80 | } 81 | 82 | // todo use constant for channel id 83 | m.loop.Request(100, evws.Connect{uri, headers}, func(err error, data interface{}) { 84 | if err != nil { 85 | L.CallByParam(lua.P{ 86 | Fn: cb, 87 | NRet: 0, 88 | Protect: false, 89 | }, lua.LString(err.Error()), lua.LNil) 90 | return 91 | } 92 | 93 | if c, ok := data.(*ws.Connection); ok { 94 | c.InitIOWorkers() 95 | conn := NewConn(c, m.loop) 96 | L.CallByParam(lua.P{ 97 | Fn: cb, 98 | NRet: 0, 99 | Protect: false, 100 | }, lua.LNil, conn.ToTable(L)) 101 | return 102 | } 103 | 104 | panic("ws: unexpected data in request callback") 105 | }) 106 | 107 | return 0 108 | })) 109 | 110 | L.Push(mod) 111 | return 1 112 | } 113 | } 114 | -------------------------------------------------------------------------------- /ev/ws/client.go: -------------------------------------------------------------------------------- 1 | package ws 2 | 3 | import ( 4 | "fmt" 5 | "github.com/gobwas/gws/ev" 6 | "github.com/gobwas/gws/ws" 7 | "net/http" 8 | "sync" 9 | "sync/atomic" 10 | ) 11 | 12 | type Connect struct { 13 | Url string 14 | Headers http.Header 15 | } 16 | 17 | type Send struct { 18 | Conn *ws.Connection 19 | Message ws.MessageRaw 20 | } 21 | 22 | type Receive struct { 23 | conn *ws.Connection 24 | } 25 | 26 | func NewReceive(c *ws.Connection) *Receive { 27 | return &Receive{ 28 | conn: c, 29 | } 30 | } 31 | 32 | type ClientHandler struct { 33 | mu sync.Mutex 34 | pending int32 35 | loops int32 36 | stop chan struct{} 37 | } 38 | 39 | func NewClientHandler() *ClientHandler { 40 | return &ClientHandler{ 41 | stop: make(chan struct{}), 42 | } 43 | } 44 | 45 | func (h *ClientHandler) Init(*ev.Loop) error { 46 | if atomic.SwapInt32(&h.loops, 1) >= 1 { 47 | return fmt.Errorf("ws handler could be registered only in one loop") 48 | } 49 | return nil 50 | } 51 | 52 | func (h *ClientHandler) Handle(loop *ev.Loop, data interface{}, cb ev.Callback) error { 53 | switch v := data.(type) { 54 | case Connect: 55 | h.doConnect(loop, v, cb) 56 | 57 | case Send: 58 | h.doSend(loop, v, cb) 59 | 60 | case *Receive: 61 | h.doReceive(loop, v, cb) 62 | 63 | default: 64 | return fmt.Errorf("unknown request format to ws handler: %s", data) 65 | } 66 | 67 | return nil 68 | } 69 | 70 | func (h *ClientHandler) Stop() { 71 | close(h.stop) 72 | } 73 | 74 | func (h *ClientHandler) IsActive(loop *ev.Loop) bool { 75 | return atomic.LoadInt32(&h.pending) > 0 76 | } 77 | 78 | func (h *ClientHandler) doConnect(loop *ev.Loop, req Connect, cb ev.Callback) { 79 | atomic.AddInt32(&h.pending, 1) 80 | go func() { 81 | conn, _, err := ws.GetConn(req.Url, req.Headers) 82 | if err != nil { 83 | loop.Call(func() { 84 | cb(err, nil) 85 | }) 86 | } else { 87 | loop.Call(func() { 88 | cb(nil, ws.NewConnection(conn)) 89 | }) 90 | } 91 | atomic.AddInt32(&h.pending, -1) 92 | }() 93 | } 94 | 95 | func (h *ClientHandler) doSend(loop *ev.Loop, req Send, cb ev.Callback) { 96 | atomic.AddInt32(&h.pending, 1) 97 | go func() { 98 | err := req.Conn.Send(req.Message) 99 | if err != nil { 100 | loop.Call(func() { 101 | cb(err, nil) 102 | }) 103 | } else { 104 | loop.Call(func() { 105 | cb(nil, nil) 106 | }) 107 | } 108 | atomic.AddInt32(&h.pending, -1) 109 | }() 110 | } 111 | 112 | func (h *ClientHandler) doReceive(loop *ev.Loop, req *Receive, cb ev.Callback) { 113 | atomic.AddInt32(&h.pending, 1) 114 | go func() { 115 | defer atomic.AddInt32(&h.pending, -1) 116 | for { 117 | select { 118 | case <-h.stop: 119 | return 120 | 121 | case <-req.conn.Done(): 122 | return 123 | 124 | case envelope := <-req.conn.ReceiveAsync(): 125 | msg, err := envelope.Message, envelope.Error 126 | if err != nil { 127 | loop.Call(func() { 128 | cb(err, nil) 129 | }) 130 | return 131 | } 132 | 133 | loop.Call(func() { 134 | cb(nil, string(msg.Data)) 135 | }) 136 | } 137 | } 138 | }() 139 | } 140 | -------------------------------------------------------------------------------- /lua/mod/time/time.go: -------------------------------------------------------------------------------- 1 | package time 2 | 3 | import ( 4 | "github.com/gobwas/gws/ev" 5 | "github.com/yuin/gopher-lua" 6 | "time" 7 | ) 8 | 9 | type durationKind int 10 | 11 | const ( 12 | durationMicroseconds durationKind = iota 13 | durationMilliseconds 14 | durationSeconds 15 | ) 16 | 17 | type Mod struct { 18 | initTime time.Time 19 | loop *ev.Loop 20 | timers map[uint32]*ev.Timer 21 | timeoutCounter uint32 22 | } 23 | 24 | func New(loop *ev.Loop) *Mod { 25 | return &Mod{ 26 | initTime: time.Now(), 27 | loop: loop, 28 | timers: make(map[uint32]*ev.Timer), 29 | } 30 | } 31 | 32 | func (m *Mod) Exports() lua.LGFunction { 33 | return func(L *lua.LState) int { 34 | mod := L.NewTable() 35 | 36 | L.SetField(mod, "us", lua.LNumber(durationMicroseconds)) 37 | L.SetField(mod, "ms", lua.LNumber(durationMilliseconds)) 38 | L.SetField(mod, "s", lua.LNumber(durationSeconds)) 39 | 40 | mod.RawSetString("now", L.NewClosure(func(L *lua.LState) int { 41 | now := inPrecision(time.Since(m.initTime), durationKind(L.ToNumber(1))) 42 | L.Push(lua.LNumber(now)) 43 | return 1 44 | })) 45 | 46 | mod.RawSetString("sleep", L.NewClosure(func(L *lua.LState) int { 47 | arg := L.ToString(1) 48 | duration, err := time.ParseDuration(arg) 49 | if err != nil { 50 | L.Push(lua.LString(err.Error())) 51 | return 1 52 | } 53 | time.Sleep(duration) 54 | return 0 55 | })) 56 | 57 | mod.RawSetString("setTimeout", L.NewClosure(func(L *lua.LState) int { 58 | tm := L.ToNumber(1) 59 | cb := L.ToFunction(2) 60 | 61 | m.timeoutCounter++ 62 | timeout := m.loop.Timeout(time.Duration(tm)*time.Millisecond, false, func() { 63 | L.CallByParam(lua.P{ 64 | Fn: cb, 65 | NRet: 0, 66 | Protect: false, 67 | }) 68 | }) 69 | m.timers[m.timeoutCounter] = timeout 70 | 71 | L.Push(lua.LNumber(m.timeoutCounter)) 72 | 73 | return 1 74 | })) 75 | 76 | mod.RawSetString("setInterval", L.NewClosure(func(L *lua.LState) int { 77 | tm := L.ToNumber(1) 78 | cb := L.ToFunction(2) 79 | 80 | m.timeoutCounter++ 81 | timeout := m.loop.Timeout(time.Duration(tm)*time.Millisecond, true, func() { 82 | L.CallByParam(lua.P{ 83 | Fn: cb, 84 | NRet: 0, 85 | Protect: false, 86 | }) 87 | }) 88 | m.timers[m.timeoutCounter] = timeout 89 | 90 | L.Push(lua.LNumber(m.timeoutCounter)) 91 | 92 | return 1 93 | })) 94 | 95 | mod.RawSetString("unsetTimer", L.NewClosure(func(L *lua.LState) int { 96 | id := L.ToNumber(1) 97 | timeout, ok := m.timers[uint32(id)] 98 | if !ok { 99 | L.Push(lua.LString("unknown timeout")) 100 | return 1 101 | } 102 | 103 | timeout.Stop() 104 | 105 | return 0 106 | })) 107 | 108 | L.Push(mod) 109 | return 1 110 | } 111 | } 112 | 113 | func inPrecision(dur time.Duration, p durationKind) float64 { 114 | switch p { 115 | case durationMicroseconds: 116 | return dur.Seconds() * 1000000 117 | case durationMilliseconds: 118 | return dur.Seconds() * 1000 119 | case durationSeconds: 120 | return dur.Seconds() 121 | 122 | default: 123 | return float64(dur) 124 | } 125 | } 126 | -------------------------------------------------------------------------------- /config/config.go: -------------------------------------------------------------------------------- 1 | // Package config brings common configuration flags and utils. 2 | package config 3 | 4 | import ( 5 | "flag" 6 | "fmt" 7 | "github.com/gobwas/gws/util/headers" 8 | headersUtil "github.com/gobwas/gws/util/headers" 9 | "net/http" 10 | "net/url" 11 | "strings" 12 | "time" 13 | ) 14 | 15 | var Verbose bool 16 | var HeaderList *headerList 17 | var Addr string 18 | var URI string 19 | var Stat time.Duration 20 | 21 | func init() { 22 | HeaderList = newHeaderList() 23 | // BoolVar and StringVar are used here just for reading them 24 | // from other packages with pure common.{Verbose|Headers} (without *) 25 | flag.BoolVar(&Verbose, "verbose", false, "verbose output") 26 | flag.StringVar(&Addr, "listen", ":3000", "address to listen") 27 | flag.StringVar(&URI, "url", ":3000", "address to connect") 28 | flag.DurationVar(&Stat, "statd", time.Second, "server statistics dump interval") 29 | flag.Var(HeaderList, "header", fmt.Sprintf("allows to specify list of headers to be passed during handshake (both in client or server)\n\tformat:\n\t\t{ key %s value }", headers.AssignmentOperator)) 30 | flag.Var(HeaderList, "H", fmt.Sprintf("allows to specify list of headers to be passed during handshake (both in client or server)\n\tformat:\n\t\t{ key %s value }", headers.AssignmentOperator)) 31 | } 32 | 33 | type headerList struct { 34 | list http.Header 35 | } 36 | 37 | func newHeaderList() *headerList { 38 | return &headerList{ 39 | list: make(http.Header), 40 | } 41 | } 42 | 43 | func (h *headerList) Set(s string) error { 44 | k, v, err := headersUtil.ParseOne(s) 45 | if err != nil { 46 | return err 47 | } 48 | h.list.Add(k, v) 49 | return nil 50 | } 51 | 52 | func (h *headerList) String() string { 53 | return fmt.Sprintf("%v", h.list) 54 | } 55 | 56 | const headerOrigin = "Origin" 57 | 58 | type Config struct { 59 | Addr string 60 | URI string 61 | Headers http.Header 62 | StatDump time.Duration 63 | } 64 | 65 | func Parse() (c Config, err error) { 66 | headers := HeaderList.list 67 | uri, err := parseURL(URI) 68 | if err != nil { 69 | return 70 | } 71 | 72 | c = Config{ 73 | Addr: Addr, 74 | URI: uri.String(), 75 | Headers: fillOriginHeader(headers, uri), 76 | StatDump: Stat, 77 | } 78 | 79 | return 80 | } 81 | 82 | func parseURL(rawURL string) (*url.URL, error) { 83 | // prevent false error on parsing url 84 | if strings.Index(rawURL, "://") == -1 { 85 | rawURL = fmt.Sprintf("ws://%s", rawURL) 86 | } 87 | uri, err := url.Parse(rawURL) 88 | if err != nil { 89 | return nil, err 90 | } 91 | if uri.Scheme == "" { 92 | uri.Scheme = "ws" 93 | } 94 | 95 | return uri, nil 96 | } 97 | 98 | func fillOriginHeader(headers http.Header, uri *url.URL) http.Header { 99 | // by default, set the same origin 100 | // to avoid same origin policy check on connections 101 | if headers.Get(headerOrigin) == "" { 102 | var s string 103 | switch uri.Scheme { 104 | case "wss": 105 | s = "https" 106 | default: 107 | s = "http" 108 | } 109 | orig := url.URL{ 110 | Scheme: s, 111 | Host: uri.Host, 112 | } 113 | headers.Set(headerOrigin, orig.String()) 114 | } 115 | 116 | return headers 117 | } 118 | -------------------------------------------------------------------------------- /stat/results/sorter/sorter.go: -------------------------------------------------------------------------------- 1 | package sorter 2 | 3 | import ( 4 | "fmt" 5 | "github.com/gobwas/gws/stat/results" 6 | "sort" 7 | "strings" 8 | ) 9 | 10 | type Sorter struct { 11 | comparators []Comparator 12 | results []results.Result 13 | } 14 | 15 | func New(s ...Comparator) *Sorter { 16 | return &Sorter{ 17 | comparators: s, 18 | } 19 | } 20 | 21 | func (l *Sorter) Sort(lines []results.Result) []results.Result { 22 | l.results = lines 23 | sort.Sort(l) 24 | return l.results 25 | } 26 | 27 | func (l *Sorter) Len() int { 28 | return len(l.results) 29 | } 30 | 31 | func (l *Sorter) Swap(i, j int) { 32 | l.results[i], l.results[j] = l.results[j], l.results[i] 33 | } 34 | 35 | func (l *Sorter) Less(a, b int) bool { 36 | lineA, lineB := l.results[a], l.results[b] 37 | for _, c := range l.comparators { 38 | switch c.compare(lineA, lineB) { 39 | case -1: 40 | return true 41 | case 1: 42 | return false 43 | } 44 | } 45 | return false 46 | } 47 | 48 | type Comparator interface { 49 | compare(a, b results.Result) int 50 | } 51 | 52 | type TagsComparator struct { 53 | TagNamesOrder []string 54 | } 55 | 56 | func (s TagsComparator) compare(a, b results.Result) int { 57 | return compareTags(s.TagNamesOrder, a.Tags, b.Tags) 58 | } 59 | 60 | type MetaComparator struct { 61 | MetaNamesOrder []string 62 | } 63 | 64 | func (s MetaComparator) compare(a, b results.Result) int { 65 | return compareMeta(s.MetaNamesOrder, a.Meta, b.Meta) 66 | } 67 | 68 | type CaptionComparator struct { 69 | } 70 | 71 | func (s CaptionComparator) compare(a, b results.Result) int { 72 | return strings.Compare(a.Name, b.Name) 73 | } 74 | 75 | type KindComparator struct{} 76 | 77 | func (s KindComparator) compare(a, b results.Result) int { 78 | return strings.Compare(a.Kind, b.Kind) 79 | } 80 | 81 | type ValueComparator struct { 82 | } 83 | 84 | func (s ValueComparator) compare(a, b results.Result) int { 85 | if a.Value < b.Value { 86 | return -1 87 | } 88 | if a.Value > b.Value { 89 | return +1 90 | } 91 | return 0 92 | } 93 | 94 | func compareTags(priority []string, a, b map[string]string) int { 95 | countTagsA, countTagsB := len(a), len(b) 96 | if countTagsA > countTagsB { 97 | return -1 98 | } 99 | if countTagsA < countTagsB { 100 | return 1 101 | } 102 | for _, tag := range priority { 103 | valueA, hasA := a[tag] 104 | valueB, hasB := b[tag] 105 | if hasA && !hasB { 106 | return -1 107 | } 108 | if hasB && !hasA { 109 | return 1 110 | } 111 | if cmp := strings.Compare(valueA, valueB); cmp != 0 { 112 | return cmp 113 | } 114 | } 115 | return 0 116 | } 117 | 118 | func compareMeta(priority []string, a, b map[string]interface{}) int { 119 | countFieldsA, countFieldsB := len(a), len(b) 120 | if countFieldsA > countFieldsB { 121 | return -1 122 | } 123 | if countFieldsA < countFieldsB { 124 | return 1 125 | } 126 | for _, key := range priority { 127 | valueA, hasA := a[key] 128 | valueB, hasB := b[key] 129 | if hasA && !hasB { 130 | return -1 131 | } 132 | if hasB && !hasA { 133 | return 1 134 | } 135 | if hasA && hasB { 136 | if cmp := strings.Compare(fmt.Sprintf("%v", valueA), fmt.Sprintf("%v", valueB)); cmp != 0 { 137 | return cmp 138 | } 139 | } 140 | } 141 | return 0 142 | } 143 | -------------------------------------------------------------------------------- /stat/stat.go: -------------------------------------------------------------------------------- 1 | package stat 2 | 3 | import ( 4 | "fmt" 5 | "github.com/gobwas/gws/stat/results" 6 | "github.com/gobwas/gws/stat/results/report" 7 | "sync" 8 | ) 9 | 10 | type Counter interface { 11 | Add(v float64) 12 | Flush() float64 13 | Kind() string 14 | } 15 | 16 | type CounterFactory func() Counter 17 | 18 | type Config struct { 19 | Factory CounterFactory 20 | Meta map[string]interface{} 21 | } 22 | 23 | type Instance struct { 24 | Counter Counter 25 | Meta map[string]interface{} 26 | } 27 | 28 | type Metric struct { 29 | Instances []Instance 30 | Tags map[string]string 31 | } 32 | 33 | func (i Metric) tagsEqual(tags map[string]string) bool { 34 | if len(tags) != len(i.Tags) { 35 | return false 36 | } 37 | for k, v := range i.Tags { 38 | if tags[k] != v { 39 | return false 40 | } 41 | } 42 | return true 43 | } 44 | 45 | type Statistics struct { 46 | mu sync.Mutex 47 | configs map[string][]*Config 48 | Metrics map[string][]*Metric 49 | } 50 | 51 | func New() *Statistics { 52 | return &Statistics{ 53 | configs: make(map[string][]*Config), 54 | Metrics: make(map[string][]*Metric), 55 | } 56 | } 57 | 58 | func (s *Statistics) Pretty() string { 59 | s.mu.Lock() 60 | defer s.mu.Unlock() 61 | 62 | report := report.New() 63 | for name, metrics := range s.Metrics { 64 | for _, metric := range metrics { 65 | for _, instance := range metric.Instances { 66 | report.AddResult(results.Result{ 67 | Name: name, 68 | Kind: instance.Counter.Kind(), 69 | Value: instance.Counter.Flush(), 70 | Tags: metric.Tags, 71 | Meta: instance.Meta, 72 | }) 73 | } 74 | } 75 | } 76 | 77 | return report.String() 78 | } 79 | 80 | func (s *Statistics) New(name string) (err error) { 81 | s.mu.Lock() 82 | { 83 | if _, ok := s.configs[name]; ok { 84 | err = fmt.Errorf("metric %q already exists", name) 85 | } else { 86 | s.configs[name] = nil 87 | } 88 | } 89 | s.mu.Unlock() 90 | return 91 | } 92 | 93 | func (s *Statistics) Setup(name string, config Config) (err error) { 94 | s.mu.Lock() 95 | { 96 | if _, ok := s.configs[name]; !ok { 97 | err = fmt.Errorf("metric %q is not exists", name) 98 | } else { 99 | s.configs[name] = append(s.configs[name], &config) 100 | } 101 | } 102 | s.mu.Unlock() 103 | return 104 | } 105 | 106 | //todo tag struct 107 | func (s *Statistics) Increment(name string, value float64, tags map[string]string) (err error) { 108 | s.mu.Lock() 109 | { 110 | configs, ok := s.configs[name] 111 | if !ok { 112 | err = fmt.Errorf("metric %q has no config", name) 113 | s.mu.Unlock() 114 | return 115 | } 116 | 117 | for _, metric := range s.Metrics[name] { 118 | if metric.tagsEqual(tags) { 119 | for _, instance := range metric.Instances { 120 | instance.Counter.Add(value) 121 | } 122 | s.mu.Unlock() 123 | return 124 | } 125 | } 126 | 127 | metric := &Metric{Tags: tags} 128 | for _, config := range configs { 129 | instance := Instance{ 130 | Meta: config.Meta, 131 | Counter: config.Factory(), 132 | } 133 | metric.Instances = append(metric.Instances, instance) 134 | instance.Counter.Add(value) 135 | } 136 | s.Metrics[name] = append(s.Metrics[name], metric) 137 | } 138 | s.mu.Unlock() 139 | return 140 | } 141 | -------------------------------------------------------------------------------- /lua/mod/ws/conn.go: -------------------------------------------------------------------------------- 1 | package ws 2 | 3 | import ( 4 | "github.com/gobwas/gws/ev" 5 | evws "github.com/gobwas/gws/ev/ws" 6 | "github.com/gobwas/gws/lua/mod" 7 | "github.com/gobwas/gws/ws" 8 | "github.com/yuin/gopher-lua" 9 | ) 10 | 11 | type Conn struct { 12 | conn *ws.Connection 13 | loop *ev.Loop 14 | emitter *mod.Emitter 15 | 16 | receive *evws.Receive 17 | listener []ln 18 | } 19 | 20 | func NewConn(c *ws.Connection, l *ev.Loop) *Conn { 21 | return &Conn{ 22 | emitter: mod.NewEmitter(), 23 | conn: c, 24 | loop: l, 25 | } 26 | } 27 | 28 | func (c *Conn) Emit(name string, args ...interface{}) { 29 | c.loop.Call(func() { 30 | c.emitter.Emit(name, args...) 31 | }) 32 | } 33 | 34 | func (c *Conn) ToTable(L *lua.LState) *lua.LTable { 35 | table := L.NewTable() 36 | 37 | table.RawSetString("send", L.NewClosure(func(L *lua.LState) int { 38 | str := L.ToString(1) 39 | msg := ws.MessageRaw{ws.TextMessage, []byte(str)} 40 | 41 | cb := L.ToFunction(2) 42 | if cb == nil { // if there is no callback - call is synchronous 43 | if err := c.conn.Send(msg); err != nil { 44 | L.Push(lua.LString(err.Error())) 45 | return 1 46 | } 47 | return 0 48 | } 49 | 50 | c.loop.Request(100, evws.Send{c.conn, msg}, func(err error, _ interface{}) { 51 | var e lua.LValue 52 | if err != nil { 53 | e = lua.LString(err.Error()) 54 | } else { 55 | e = lua.LNil 56 | } 57 | L.CallByParam(lua.P{ 58 | Fn: cb, 59 | NRet: 0, 60 | Protect: false, 61 | }, e) 62 | }) 63 | return 0 64 | })) 65 | 66 | table.RawSetString("listen", L.NewClosure(func(L *lua.LState) int { 67 | cb := L.ToFunction(1) 68 | c.listener = append(c.listener, ln{cb, L}) 69 | 70 | if c.receive == nil { 71 | c.receive = evws.NewReceive(c.conn) 72 | c.loop.Request(100, c.receive, func(err error, msg interface{}) { 73 | for _, ln := range c.listener { 74 | if err != nil { 75 | ln.state.CallByParam(lua.P{ 76 | Fn: cb, 77 | NRet: 0, 78 | Protect: false, 79 | }, lua.LString(err.Error()), lua.LNil) 80 | } else { 81 | ln.state.CallByParam(lua.P{ 82 | Fn: cb, 83 | NRet: 0, 84 | Protect: false, 85 | }, lua.LNil, lua.LString(msg.(string))) 86 | } 87 | } 88 | }) 89 | } 90 | 91 | return 0 92 | })) 93 | 94 | table.RawSetString("receive", L.NewClosure(func(L *lua.LState) int { 95 | if c.receive != nil { 96 | L.Push(lua.LString("could not receive synchronous: there are already registered listeners")) 97 | return 1 98 | } 99 | 100 | msg, err := c.conn.Receive() 101 | if err != nil { 102 | L.Push(lua.LNil) 103 | L.Push(lua.LString(err.Error())) 104 | return 2 105 | } else { 106 | L.Push(lua.LString(msg.Data)) 107 | L.Push(lua.LNil) 108 | return 2 109 | } 110 | })) 111 | 112 | table.RawSetString("close", L.NewClosure(func(L *lua.LState) int { 113 | if err := c.conn.Close(); err != nil { 114 | L.Push(lua.LString(err.Error())) 115 | return 1 116 | } 117 | 118 | c.Emit("close") 119 | return 0 120 | })) 121 | 122 | table.RawSetString("on", c.emitter.ExportOn(L)) 123 | table.RawSetString("off", c.emitter.ExportOff(L)) 124 | 125 | return table 126 | } 127 | 128 | type ln struct { 129 | cb *lua.LFunction 130 | state *lua.LState 131 | } 132 | -------------------------------------------------------------------------------- /ws/server.go: -------------------------------------------------------------------------------- 1 | package ws 2 | 3 | import ( 4 | "crypto/tls" 5 | "fmt" 6 | "github.com/gorilla/websocket" 7 | "net" 8 | "net/http" 9 | "sync" 10 | "time" 11 | ) 12 | 13 | type Handler interface { 14 | Handle(*websocket.Conn, error) 15 | } 16 | 17 | type HandlerFunc func(*websocket.Conn, error) 18 | 19 | func (h HandlerFunc) Handle(c *websocket.Conn, e error) { 20 | h(c, e) 21 | } 22 | 23 | type ServerConfig struct { 24 | Addr string 25 | Key string 26 | Cert string 27 | Origin string 28 | Headers http.Header 29 | } 30 | 31 | type Server struct { 32 | mu sync.Mutex 33 | listening bool 34 | config ServerConfig 35 | deferreds []func() 36 | handlers []Handler 37 | conns chan conn 38 | } 39 | 40 | func NewServer(cfg ServerConfig) *Server { 41 | return &Server{ 42 | config: cfg, 43 | conns: make(chan conn), 44 | } 45 | } 46 | 47 | func (s *Server) Handle(h Handler) { 48 | s.mu.Lock() 49 | defer s.mu.Unlock() 50 | 51 | s.handlers = append(s.handlers, h) 52 | go func() { 53 | for w := range s.conns { 54 | h.Handle(w.conn, w.err) 55 | } 56 | }() 57 | } 58 | 59 | func (s *Server) Defer(dfd func()) { 60 | s.mu.Lock() 61 | defer s.mu.Unlock() 62 | s.deferreds = append(s.deferreds, dfd) 63 | } 64 | 65 | func (s *Server) Listen(done chan struct{}) { 66 | go func() { 67 | upgrade := GetUpgrader(UpgradeConfig{ 68 | Origin: s.config.Origin, 69 | Headers: s.config.Headers, 70 | }) 71 | 72 | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 73 | c, err := upgrade(w, r) 74 | s.conns <- conn{c, err} 75 | }) 76 | 77 | var ln net.Listener 78 | var err error 79 | if s.config.Cert != "" || s.config.Key != "" { 80 | ln, err = getTLSListener(done, s.config.Addr, s.config.Cert, s.config.Key) 81 | } else { 82 | ln, err = getListener(done, s.config.Addr) 83 | } 84 | 85 | if err == nil { 86 | err = http.Serve(ln, handler) 87 | } 88 | 89 | s.mu.Lock() 90 | defer s.mu.Unlock() 91 | for _, dfd := range s.deferreds { 92 | dfd() 93 | } 94 | close(s.conns) 95 | for _, h := range s.handlers { 96 | h.Handle(nil, err) 97 | } 98 | }() 99 | } 100 | 101 | type conn struct { 102 | conn *websocket.Conn 103 | err error 104 | } 105 | 106 | type tcpStoppableListener struct { 107 | *net.TCPListener 108 | stop chan struct{} 109 | } 110 | 111 | func (ln *tcpStoppableListener) Accept() (c net.Conn, err error) { 112 | for { 113 | ln.SetDeadline(time.Now().Add(time.Second)) 114 | select { 115 | case <-ln.stop: 116 | ln.Close() 117 | return nil, fmt.Errorf("listener has been stopped!") 118 | 119 | default: 120 | c, err = ln.AcceptTCP() 121 | if err != nil { 122 | if ne, ok := err.(net.Error); ok && ne.Temporary() && ne.Timeout() { 123 | continue 124 | } 125 | } 126 | return 127 | } 128 | } 129 | } 130 | 131 | func getListener(done chan struct{}, addr string) (net.Listener, error) { 132 | ln, err := net.Listen("tcp", addr) 133 | if err != nil { 134 | return nil, err 135 | } 136 | 137 | return &tcpStoppableListener{ln.(*net.TCPListener), done}, nil 138 | } 139 | 140 | func getTLSListener(done chan struct{}, addr, cert, key string) (net.Listener, error) { 141 | config := &tls.Config{ 142 | NextProtos: []string{"http/1.1"}, 143 | Certificates: make([]tls.Certificate, 1), 144 | } 145 | var err error 146 | config.Certificates[0], err = tls.LoadX509KeyPair(cert, key) 147 | if err != nil { 148 | return nil, err 149 | } 150 | 151 | ln, err := net.Listen("tcp", addr) 152 | if err != nil { 153 | return nil, err 154 | } 155 | 156 | return tls.NewListener(&tcpStoppableListener{ln.(*net.TCPListener), done}, config), nil 157 | } 158 | -------------------------------------------------------------------------------- /lua/mod/mod.go: -------------------------------------------------------------------------------- 1 | package mod 2 | 3 | import ( 4 | "github.com/gobwas/gws/lua/util" 5 | "github.com/yuin/gopher-lua" 6 | ) 7 | 8 | type Module interface { 9 | Exports() lua.LGFunction 10 | } 11 | 12 | type Storage struct { 13 | data map[string]interface{} 14 | } 15 | 16 | func NewStorage() *Storage { 17 | return &Storage{ 18 | data: make(map[string]interface{}), 19 | } 20 | } 21 | 22 | func (s *Storage) Set(key string, value interface{}) { 23 | s.data[key] = value 24 | return 25 | } 26 | 27 | func (s *Storage) Get(key string) (result interface{}) { 28 | result = s.data[key] 29 | return 30 | } 31 | 32 | func (s *Storage) ExportSet(L *lua.LState) *lua.LFunction { 33 | return L.NewClosure(func(L *lua.LState) int { 34 | key := L.ToString(1) 35 | value := L.Get(3) 36 | s.Set(key, value) 37 | return 0 38 | }) 39 | } 40 | 41 | func (s *Storage) ExportGet(L *lua.LState) *lua.LFunction { 42 | return L.NewClosure(func(L *lua.LState) int { 43 | key := L.ToString(1) 44 | switch v := s.Get(key).(type) { 45 | case map[string]string: 46 | L.Push(util.MapOfStringToTable(L, v)) 47 | case map[string]interface{}: 48 | L.Push(util.MapOfInterfaceToTable(L, v)) 49 | case lua.LValue: 50 | L.Push(v) 51 | case string: 52 | L.Push(lua.LString(v)) 53 | case int: 54 | L.Push(lua.LNumber(v)) 55 | case uint: 56 | L.Push(lua.LNumber(v)) 57 | case float64: 58 | L.Push(lua.LNumber(v)) 59 | default: 60 | L.Push(lua.LNil) 61 | } 62 | return 1 63 | }) 64 | } 65 | 66 | type Emitter struct { 67 | seq uint32 68 | callbacks map[string][]*callback 69 | registry map[desc]uint32 70 | } 71 | 72 | func NewEmitter() *Emitter { 73 | return &Emitter{ 74 | callbacks: make(map[string][]*callback), 75 | registry: make(map[desc]uint32), 76 | } 77 | } 78 | 79 | func (e *Emitter) Emit(name string, args ...interface{}) { 80 | if len(e.callbacks[name]) == 0 { 81 | return 82 | } 83 | for _, callback := range e.callbacks[name] { 84 | callback.cb(args...) 85 | } 86 | } 87 | 88 | func (e *Emitter) On(event string, cb func(...interface{})) uint32 { 89 | e.seq++ 90 | e.callbacks[event] = append(e.callbacks[event], &callback{e.seq, cb}) 91 | return e.seq 92 | } 93 | 94 | func (e *Emitter) Off(id uint32) { 95 | for evt, callbacks := range e.callbacks { 96 | for i, cb := range callbacks { 97 | if cb.id == id { 98 | copy(callbacks[i:], callbacks[i+1:]) 99 | last := len(callbacks) + 1 100 | callbacks[last] = nil 101 | e.callbacks[evt] = callbacks[:last] 102 | return 103 | } 104 | } 105 | } 106 | } 107 | 108 | func (e *Emitter) ExportOn(L *lua.LState) *lua.LFunction { 109 | return L.NewClosure(func(L *lua.LState) int { 110 | name := L.ToString(1) 111 | cb := L.ToFunction(2) 112 | e.registry[desc{name, cb}] = e.On(name, func(args ...interface{}) { 113 | var callArgs []lua.LValue 114 | for _, arg := range args { 115 | switch v := arg.(type) { 116 | case string: 117 | callArgs = append(callArgs, lua.LString(v)) 118 | case int: 119 | callArgs = append(callArgs, lua.LNumber(v)) 120 | case uint: 121 | callArgs = append(callArgs, lua.LNumber(v)) 122 | case float64: 123 | callArgs = append(callArgs, lua.LNumber(v)) 124 | default: 125 | // 126 | } 127 | } 128 | L.CallByParam(lua.P{ 129 | Fn: cb, 130 | NRet: 0, 131 | Protect: false, 132 | }, callArgs...) 133 | }) 134 | return 0 135 | }) 136 | } 137 | 138 | func (e *Emitter) ExportOff(L *lua.LState) *lua.LFunction { 139 | return L.NewClosure(func(L *lua.LState) int { 140 | name := L.ToString(1) 141 | cb := L.ToFunction(2) 142 | if id, ok := e.registry[desc{name, cb}]; ok { 143 | e.Off(id) 144 | } 145 | return 0 146 | }) 147 | } 148 | 149 | type callback struct { 150 | id uint32 151 | cb func(...interface{}) 152 | } 153 | 154 | type desc struct { 155 | event string 156 | fn *lua.LFunction 157 | } 158 | -------------------------------------------------------------------------------- /stat/results/report/report.go: -------------------------------------------------------------------------------- 1 | package report 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "github.com/gobwas/gws/stat/results" 7 | "github.com/gobwas/gws/stat/results/sorter" 8 | "io" 9 | "sort" 10 | "strings" 11 | ) 12 | 13 | const ( 14 | titleCaption = "counter" 15 | titleKind = "kind" 16 | titleValue = "value" 17 | ) 18 | 19 | type Report struct { 20 | results []results.Result 21 | fieldSize map[string]int 22 | meta map[string]bool 23 | tags map[string]bool 24 | tab string 25 | } 26 | 27 | func New() *Report { 28 | return &Report{ 29 | fieldSize: make(map[string]int), 30 | meta: make(map[string]bool), 31 | tags: make(map[string]bool), 32 | tab: " ", 33 | } 34 | } 35 | 36 | func (p *Report) AddResult(result results.Result) { 37 | p.results = append(p.results, result) 38 | p.updateFieldsSizes(result) 39 | } 40 | 41 | func (p *Report) String() string { 42 | tags := p.getTagNamesOrder() 43 | meta := p.getMetaNamesOrder() 44 | 45 | var fields []string 46 | fields = append(fields, titleCaption) 47 | fields = append(fields, tags...) 48 | fields = append(fields, titleValue) 49 | fields = append(fields, titleKind) 50 | fields = append(fields, meta...) 51 | 52 | buf := &bytes.Buffer{} 53 | p.printHeader(buf, fields) 54 | for _, result := range p.getResultsOrdered(tags, meta) { 55 | var line []pair 56 | line = append(line, pair{titleCaption, result.Name}) 57 | line = append(line, resultTagPairs(tags, result)...) 58 | line = append(line, pair{titleValue, valueToString(result.Value)}) 59 | line = append(line, pair{titleKind, result.Kind}) 60 | line = append(line, resultMetaPairs(meta, result)...) 61 | p.printLine(buf, line) 62 | } 63 | 64 | buf.WriteByte('\n') 65 | 66 | return buf.String() 67 | } 68 | 69 | func (p *Report) getResultsOrdered(tagsOrder, metaOrder []string) []results.Result { 70 | orderer := sorter.New( 71 | &sorter.CaptionComparator{}, 72 | &sorter.TagsComparator{tagsOrder}, 73 | &sorter.KindComparator{}, 74 | &sorter.ValueComparator{}, 75 | ) 76 | return orderer.Sort(p.results) 77 | } 78 | 79 | func (p *Report) updateFieldsSizes(result results.Result) { 80 | p.updateFieldSize(titleCaption, result.Name) 81 | p.updateFieldSize(titleKind, result.Kind) 82 | p.updateFieldSize(titleValue, valueToString(result.Value)) 83 | 84 | for k, v := range result.Meta { 85 | p.updateFieldSize(k, fmt.Sprintf("%v", v)) 86 | p.meta[k] = true 87 | } 88 | 89 | for k, v := range result.Tags { 90 | p.updateFieldSize(k, v) 91 | p.tags[k] = true 92 | } 93 | } 94 | 95 | func (p *Report) updateFieldSize(name string, value string) { 96 | if cur, ok := p.fieldSize[name]; ok { 97 | if l := len(value); l > cur { 98 | p.fieldSize[name] = l 99 | } 100 | return 101 | } 102 | 103 | if ln, lv := len(name), len(value); lv > ln { 104 | p.fieldSize[name] = lv 105 | } else { 106 | p.fieldSize[name] = ln 107 | } 108 | } 109 | 110 | func (p *Report) getMetaNamesOrder() (fields []string) { 111 | fields = make([]string, 0, len(p.meta)) 112 | for k := range p.meta { 113 | fields = append(fields, k) 114 | } 115 | sort.Strings(fields) 116 | return 117 | } 118 | 119 | func (p *Report) getTagNamesOrder() (fields []string) { 120 | fields = make([]string, 0, len(p.tags)) 121 | for t := range p.tags { 122 | fields = append(fields, t) 123 | } 124 | sort.Strings(fields) 125 | return 126 | } 127 | 128 | func (p *Report) printHeader(w io.Writer, fields []string) (n int, err error) { 129 | var line []pair 130 | for _, f := range fields { 131 | line = append(line, pair{f, f}) 132 | } 133 | n, err = p.printLine(w, line) 134 | if err != nil { 135 | return 136 | } 137 | w.Write([]byte(strings.Repeat("-", n-1))) 138 | w.Write([]byte{'\n'}) 139 | return 140 | } 141 | 142 | func (p *Report) printLine(w io.Writer, line []pair) (n int, err error) { 143 | fieldsLen := len(line) 144 | if fieldsLen == 0 { 145 | return 146 | } 147 | var length int 148 | for i, pair := range line { 149 | size := p.fieldSize[pair.key] 150 | var eol string 151 | if i+1 == fieldsLen { 152 | eol = "\n" 153 | } else { 154 | eol = p.tab 155 | } 156 | 157 | n, err = fmt.Fprintf(w, "%-*s%s", size, pair.value, eol) 158 | if err != nil { 159 | return 160 | } 161 | length += n 162 | } 163 | 164 | return length, nil 165 | } 166 | 167 | func resultTagPairs(tags []string, result results.Result) (pairs []pair) { 168 | for _, key := range tags { 169 | var value string 170 | if v, ok := result.Tags[key]; ok { 171 | value = v 172 | } 173 | pairs = append(pairs, pair{key, value}) 174 | } 175 | return 176 | } 177 | 178 | func resultMetaPairs(meta []string, result results.Result) (pairs []pair) { 179 | for _, key := range meta { 180 | var value string 181 | if v, ok := result.Meta[key]; ok { 182 | value = fmt.Sprintf("%s", v) 183 | } 184 | pairs = append(pairs, pair{key, value}) 185 | } 186 | return 187 | } 188 | 189 | func valueToString(value float64) string { 190 | return fmt.Sprintf("%.3f", value) 191 | } 192 | 193 | type pair struct { 194 | key, value string 195 | } 196 | -------------------------------------------------------------------------------- /lua/lua.go: -------------------------------------------------------------------------------- 1 | package lua 2 | 3 | import ( 4 | "bytes" 5 | "flag" 6 | "fmt" 7 | "io/ioutil" 8 | "log" 9 | "os" 10 | "os/signal" 11 | "strings" 12 | "sync" 13 | "time" 14 | 15 | "github.com/gobwas/gws/bufio" 16 | "github.com/gobwas/gws/cli" 17 | "github.com/gobwas/gws/cli/color" 18 | "github.com/gobwas/gws/config" 19 | "github.com/gobwas/gws/display" 20 | "github.com/gobwas/gws/ev" 21 | evWS "github.com/gobwas/gws/ev/ws" 22 | modRuntime "github.com/gobwas/gws/lua/mod/runtime" 23 | modStat "github.com/gobwas/gws/lua/mod/stat" 24 | modTime "github.com/gobwas/gws/lua/mod/time" 25 | modWS "github.com/gobwas/gws/lua/mod/ws" 26 | "github.com/gobwas/gws/lua/script" 27 | "github.com/gobwas/gws/lua/util" 28 | "github.com/gobwas/gws/stat" 29 | ) 30 | 31 | var scriptFile = flag.String("path", "", "path to lua script") 32 | var useDisplay = flag.Bool("display", false, "use display ouput") 33 | 34 | func initRunTime(loop *ev.Loop, c config.Config) *modRuntime.Runtime { 35 | rtime := modRuntime.New(loop) 36 | rtime.Set("url", c.URI) 37 | rtime.Set("listen", c.Addr) 38 | rtime.Set("headers", util.HeadersToMap(c.Headers)) 39 | 40 | environment := os.Environ() 41 | env := make(map[string]string, len(environment)) 42 | for _, e := range environment { 43 | pair := strings.Split(e, "=") 44 | env[pair[0]] = pair[1] 45 | } 46 | rtime.Set("env", env) 47 | 48 | return rtime 49 | } 50 | 51 | func Go(c config.Config) error { 52 | var code string 53 | if script, err := ioutil.ReadFile(*scriptFile); err != nil { 54 | return err 55 | } else { 56 | code = string(script) 57 | } 58 | 59 | stats := stat.New() 60 | 61 | luaOutputBuffer := bytes.NewBuffer(make([]byte, 0, 1<<13)) 62 | luaStdout := bufio.NewWriter(luaOutputBuffer, 1<<13) 63 | 64 | systemStdout := bytes.NewBuffer(make([]byte, 0, 1024)) 65 | 66 | printer := display.NewDisplay(os.Stderr, display.Config{ 67 | TabSize: 4, 68 | Interval: time.Millisecond * 100, 69 | }) 70 | printer.Row().Col(-1, -1, func() string { 71 | return stats.Pretty() 72 | }) 73 | printer.Row().Col(256, 10, func() (str string) { 74 | luaStdout.Dump() 75 | str = luaOutputBuffer.String() 76 | luaOutputBuffer.Reset() 77 | return 78 | }) 79 | printer.Row().Col(256, 3, func() (str string) { 80 | str = systemStdout.String() 81 | return 82 | }) 83 | 84 | printer.On() 85 | defer printer.Off() 86 | defer printer.Render() 87 | 88 | cancel := make(chan struct{}) 89 | go func() { 90 | c := make(chan os.Signal, 1) 91 | signal.Notify(c, os.Interrupt) 92 | s := <-c 93 | fmt.Fprintln(systemStdout, color.Cyan(cli.PrefixTheEnd), s.String()) 94 | fmt.Fprintln(systemStdout, color.Cyan("stopping softly..")) 95 | close(cancel) 96 | s = <-c 97 | fmt.Fprintln(systemStdout, color.Red(cli.PrefixTheEnd), color.Yellow(s.String()+"x2")) 98 | fmt.Fprintln(systemStdout, color.Red("stopping hardly..")) 99 | printer.Off() 100 | os.Exit(1) 101 | }() 102 | 103 | luaScript := script.New() 104 | defer luaScript.Shutdown() 105 | 106 | luaScript.HijackOutput(bufio.NewPrefixWriter(luaStdout, color.Green("master > "))) 107 | 108 | loop := ev.NewLoop() 109 | 110 | loopServerHandler := evWS.NewServerHandler() 111 | loop.Register(evWS.NewClientHandler(), 100) 112 | loop.Register(loopServerHandler, 101) 113 | 114 | sharedStat := modStat.New(stats) 115 | 116 | var wg sync.WaitGroup 117 | var threads int 118 | rtime := initRunTime(loop, c) 119 | rtime.SetForkFn(func() error { 120 | go func(id int) { 121 | defer wg.Done() 122 | 123 | luaScript := script.New() 124 | defer luaScript.Shutdown() 125 | luaScript.HijackOutput(bufio.NewPrefixWriter(luaStdout, color.Green(fmt.Sprintf("thread %.2d > ", id)))) 126 | 127 | loop := ev.NewLoop() 128 | loop.Register(evWS.NewClientHandler(), 100) 129 | loop.Register(loopServerHandler, 101) 130 | 131 | rtime := initRunTime(loop, c) 132 | rtime.Set("id", id) 133 | 134 | luaScript.Preload("runtime", rtime) 135 | luaScript.Preload("stat", sharedStat) 136 | luaScript.Preload("time", modTime.New(loop)) 137 | luaScript.Preload("ws", modWS.New(loop)) 138 | 139 | err := luaScript.Do(code) 140 | if err != nil { 141 | log.Printf("run forked lua script error: %s", err) 142 | } 143 | 144 | loop.Run() 145 | loop.Teardown(func() { 146 | rtime.Emit("exit") 147 | }) 148 | 149 | waitLoop(cancel, loop) 150 | }(threads) 151 | 152 | wg.Add(1) 153 | threads++ 154 | 155 | return nil 156 | }) 157 | 158 | luaScript.Preload("runtime", rtime) 159 | luaScript.Preload("stat", sharedStat) 160 | luaScript.Preload("time", modTime.New(loop)) 161 | luaScript.Preload("ws", modWS.New(loop)) 162 | 163 | err := luaScript.Do(code) 164 | if err != nil { 165 | log.Printf("run lua script error: %s", err) 166 | return err 167 | } 168 | 169 | loop.Run() 170 | loop.Teardown(func() { 171 | wg.Wait() 172 | rtime.Emit("exit") 173 | }) 174 | 175 | waitLoop(cancel, loop) 176 | wg.Wait() 177 | 178 | return nil 179 | } 180 | 181 | func waitLoop(cancel chan struct{}, loop *ev.Loop) { 182 | select { 183 | case <-loop.Done(): 184 | case <-cancel: 185 | loop.Stop() 186 | shutdown := time.NewTimer(time.Second * 4) 187 | select { 188 | case <-loop.Done(): 189 | case <-shutdown.C: 190 | loop.Shutdown() 191 | } 192 | } 193 | } 194 | -------------------------------------------------------------------------------- /ws/ws.go: -------------------------------------------------------------------------------- 1 | package ws 2 | 3 | import ( 4 | "crypto/tls" 5 | "flag" 6 | "fmt" 7 | "io" 8 | "io/ioutil" 9 | "net" 10 | "net/http" 11 | "time" 12 | 13 | "github.com/gobwas/glob" 14 | "github.com/gorilla/websocket" 15 | ) 16 | 17 | var insecure = flag.Bool("insecure", false, "do not check tls certificate during dialing") 18 | var keepalive = flag.Duration("keepalive", time.Minute*60, "how long to ws connection should be alive") 19 | 20 | type Kind int 21 | 22 | const HeaderOrigin = "Origin" 23 | 24 | const ( 25 | TextMessage = 1 26 | BinaryMessage = 2 27 | CloseMessage = 8 28 | PingMessage = 9 29 | PongMessage = 10 30 | ) 31 | 32 | func (m Kind) String() string { 33 | switch m { 34 | case TextMessage: 35 | return "TextMessage" 36 | case BinaryMessage: 37 | return "BinaryMessage" 38 | case CloseMessage: 39 | return "CloseMessage" 40 | case PingMessage: 41 | return "PingMessage" 42 | case PongMessage: 43 | return "PongMessage" 44 | default: 45 | return "UnknownMessage" 46 | } 47 | } 48 | 49 | type WriteRequest struct { 50 | Message MessageRaw 51 | Result chan error 52 | } 53 | 54 | type ReceiveRequest struct { 55 | Result chan MessageAndError 56 | } 57 | 58 | type MessageRaw struct { 59 | Kind Kind 60 | Data []byte 61 | } 62 | 63 | type MessageAndError struct { 64 | Message MessageRaw 65 | Error error 66 | } 67 | 68 | type Message struct { 69 | Kind Kind 70 | Data []byte 71 | Err error 72 | } 73 | 74 | func WriteToConnFromChan(done <-chan struct{}, conn *websocket.Conn, output <-chan WriteRequest) { 75 | go func() { 76 | for { 77 | select { 78 | case <-done: 79 | return 80 | 81 | case req := <-output: 82 | err := WriteToConn(conn, req.Message.Kind, req.Message.Data) 83 | select { 84 | case <-done: 85 | case req.Result <- err: 86 | } 87 | } 88 | } 89 | }() 90 | } 91 | 92 | func ReadFromConnToChan(done <-chan struct{}, conn *websocket.Conn, ch <-chan ReceiveRequest) { 93 | go func() { 94 | for { 95 | select { 96 | case <-done: 97 | return 98 | 99 | case req := <-ch: 100 | m, err := ReadFromConn(conn) 101 | select { 102 | case <-done: 103 | case req.Result <- MessageAndError{m, err}: 104 | } 105 | } 106 | } 107 | }() 108 | } 109 | 110 | func WriteToConn(conn *websocket.Conn, t Kind, b []byte) error { 111 | writer, err := conn.NextWriter(int(t)) 112 | if err != nil { 113 | return err 114 | } 115 | 116 | _, err = writer.Write(b) 117 | if err != nil { 118 | return err 119 | } 120 | 121 | err = writer.Close() 122 | if err != nil { 123 | return err 124 | } 125 | 126 | return nil 127 | } 128 | 129 | func ReadFromConn(conn *websocket.Conn) (msg MessageRaw, err error) { 130 | t, r, err := conn.NextReader() 131 | if err != nil { 132 | return 133 | } 134 | 135 | b, err := ioutil.ReadAll(r) 136 | if err != nil { 137 | return 138 | } 139 | 140 | msg = MessageRaw{Data: b, Kind: Kind(t)} 141 | return 142 | } 143 | 144 | //todo refactor this to dedup 145 | func ReadFromConnInto(done <-chan struct{}, conn *websocket.Conn, ch chan<- Message) { 146 | go func() { 147 | closeHandler := conn.CloseHandler() 148 | defer conn.SetCloseHandler(closeHandler) 149 | conn.SetCloseHandler(func(code int, text string) error { 150 | ch <- Message{ 151 | Data: []byte(fmt.Sprintf("%d: %s", code, text)), 152 | Kind: CloseMessage, 153 | } 154 | return closeHandler(code, text) 155 | }) 156 | 157 | pingHandler := conn.PingHandler() 158 | defer conn.SetPingHandler(pingHandler) 159 | conn.SetPingHandler(func(data string) error { 160 | ch <- Message{Data: []byte(data), Kind: PingMessage} 161 | return pingHandler(data) 162 | }) 163 | 164 | pongHandler := conn.PongHandler() 165 | defer conn.SetPongHandler(pongHandler) 166 | conn.SetPongHandler(func(data string) error { 167 | ch <- Message{Data: []byte(data), Kind: PongMessage} 168 | return pongHandler(data) 169 | }) 170 | 171 | for { 172 | var msg Message 173 | t, r, err := conn.NextReader() 174 | if err != nil { 175 | msg = Message{Err: io.EOF} 176 | } else { 177 | b, err := ioutil.ReadAll(r) 178 | if err != nil { 179 | msg = Message{Err: err} 180 | } else { 181 | msg = Message{Data: b, Kind: Kind(t)} 182 | } 183 | } 184 | 185 | select { 186 | case <-done: 187 | return 188 | case ch <- msg: 189 | if msg.Err != nil { 190 | return 191 | } 192 | } 193 | } 194 | }() 195 | } 196 | 197 | func ReadAsyncFromConn(done <-chan struct{}, conn *websocket.Conn) <-chan Message { 198 | ch := make(chan Message) 199 | ReadFromConnInto(done, conn, ch) 200 | return ch 201 | } 202 | 203 | func GetConn(uri string, h http.Header) (conn *websocket.Conn, resp *http.Response, err error) { 204 | dialer := &websocket.Dialer{ 205 | NetDial: func(network, addr string) (net.Conn, error) { 206 | netDialer := &net.Dialer{ 207 | KeepAlive: *keepalive, 208 | } 209 | return netDialer.Dial(network, addr) 210 | }, 211 | } 212 | if *insecure { 213 | dialer.TLSClientConfig = &tls.Config{ 214 | InsecureSkipVerify: true, 215 | } 216 | } 217 | conn, resp, err = dialer.Dial(uri, h) 218 | return 219 | } 220 | 221 | type UpgradeConfig struct { 222 | Origin string 223 | Headers http.Header 224 | } 225 | 226 | type Upgrader func(http.ResponseWriter, *http.Request) (*websocket.Conn, error) 227 | 228 | func GetUpgrader(config UpgradeConfig) Upgrader { 229 | u := &websocket.Upgrader{} 230 | if config.Origin != "" { 231 | originChecker := glob.MustCompile(config.Origin) 232 | u.CheckOrigin = func(r *http.Request) bool { 233 | return originChecker.Match(r.Header.Get(HeaderOrigin)) 234 | } 235 | } 236 | 237 | return func(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) { 238 | return u.Upgrade(w, r, config.Headers) 239 | } 240 | } 241 | -------------------------------------------------------------------------------- /display/display.go: -------------------------------------------------------------------------------- 1 | // Package display brings tools for multiscreen output. 2 | package display 3 | 4 | import ( 5 | "bytes" 6 | "fmt" 7 | "io" 8 | "sync" 9 | "time" 10 | ) 11 | 12 | const ( 13 | cursor_hide = "\033[?25l" 14 | cursor_show = "\033[?25h" 15 | cursor_move_top = "\033[%dA" 16 | cursor_move_left = "\033[%dD" 17 | ) 18 | 19 | var ( 20 | space = []byte{' '} 21 | tab = []byte{'\t'} 22 | newl = []byte{'\n'} 23 | ) 24 | 25 | type ContentFn func() string 26 | 27 | type Column struct { 28 | Width int 29 | Height int 30 | Content ContentFn 31 | } 32 | 33 | type Row struct { 34 | mu sync.Mutex 35 | columns []Column 36 | } 37 | 38 | // Column adds new column for given row. 39 | func (r *Row) Column(z Column) *Row { 40 | r.mu.Lock() 41 | defer r.mu.Unlock() 42 | r.columns = append(r.columns, z) 43 | return r 44 | } 45 | 46 | // Col is the same as Column, but avoids creation of struct Columng by a client. 47 | func (r *Row) Col(width, height int, c ContentFn) *Row { 48 | r.mu.Lock() 49 | defer r.mu.Unlock() 50 | r.columns = append(r.columns, Column{width, height, c}) 51 | return r 52 | } 53 | 54 | // Config contains fields for display configuration. 55 | type Config struct { 56 | TabSize int // how big in spaces should be '\t' characters 57 | Interval time.Duration // interval for re-render output 58 | } 59 | 60 | // Display represents multiscreen display. 61 | type Display struct { 62 | mu sync.Mutex 63 | buf []byte 64 | dest io.Writer 65 | index int64 66 | height int 67 | width int 68 | rows []*Row 69 | done chan struct{} 70 | config Config 71 | } 72 | 73 | // NewDisplay creates new multiscreen display that renders to w. 74 | func NewDisplay(w io.Writer, c Config) *Display { 75 | if c.TabSize == 0 { 76 | c.TabSize = 4 77 | } 78 | if c.Interval == 0 { 79 | c.Interval = time.Millisecond * 100 80 | } 81 | return &Display{ 82 | dest: w, 83 | done: make(chan struct{}), 84 | config: c, 85 | } 86 | } 87 | 88 | // Row creates new row in display. 89 | func (d *Display) Row() (r *Row) { 90 | d.mu.Lock() 91 | defer d.mu.Unlock() 92 | r = &Row{} 93 | d.rows = append(d.rows, r) 94 | return 95 | } 96 | 97 | // On starts the rendering loop of display contents. 98 | func (d *Display) On() { 99 | d.dest.Write([]byte(cursor_hide)) 100 | go d.renderLoop(d.config.Interval) 101 | } 102 | 103 | // Off stops the rendering loop. 104 | func (d *Display) Off() { 105 | close(d.done) 106 | d.dest.Write([]byte(cursor_show)) 107 | } 108 | 109 | // Render renders all rows and columns registered in this Display. 110 | func (d *Display) Render() { 111 | d.mu.Lock() 112 | defer d.mu.Unlock() 113 | 114 | // clear previously rendered data 115 | if d.height > 0 { 116 | d.buf = append(d.buf, fmt.Sprintf(cursor_move_top, d.height)...) 117 | } 118 | if d.width > 0 { 119 | d.buf = append(d.buf, fmt.Sprintf(cursor_move_left, d.width)...) 120 | } 121 | 122 | var height, width int 123 | for _, row := range d.rows { 124 | var rowLines [][]byte 125 | 126 | var currentPadding int 127 | for _, col := range row.columns { 128 | data := []byte(col.Content()) 129 | data = bytes.Replace(data, tab, bytes.Repeat(space, d.config.TabSize), -1) 130 | 131 | lines := splitToLines(data, col.Width) 132 | if col.Height > 0 { 133 | if len(lines) > col.Height { 134 | // drop stale lines and keep only those that fits into column height 135 | lines = lines[len(lines)-col.Height:] 136 | } else { 137 | // if we have not reached the height of column 138 | // then fill it with empty lines 139 | for i := 0; i < col.Height-len(lines); i++ { 140 | lines = append(lines, repeat(' ', col.Width)) 141 | } 142 | } 143 | } 144 | 145 | var maxPad int 146 | for i, ln := range lines { 147 | if i == len(rowLines) { 148 | // if previous column did not filled current line 149 | // fill it with spaces 150 | rowLines = append(rowLines, repeat(' ', currentPadding)) 151 | height++ 152 | } 153 | 154 | rowLines[i] = append(rowLines[i], ' ') // write column delimiter 155 | rowLines[i] = append(rowLines[i], ln...) 156 | 157 | width = max(len(rowLines[i]), width) 158 | maxPad = max(len(rowLines[i]), maxPad) 159 | } 160 | currentPadding = maxPad 161 | } 162 | 163 | d.buf = append(d.buf, bytes.Join(rowLines, newl)...) 164 | d.buf = append(d.buf, '\n') 165 | } 166 | 167 | d.height = height 168 | d.width = width 169 | 170 | io.Copy(d.dest, bytes.NewReader(d.buf)) 171 | d.buf = d.buf[:0] 172 | } 173 | 174 | func (d *Display) renderLoop(duration time.Duration) { 175 | ticker := time.Tick(duration) 176 | for { 177 | select { 178 | case <-ticker: 179 | d.Render() 180 | case <-d.done: 181 | return 182 | } 183 | } 184 | } 185 | 186 | // splitToLines splits given data by '\n' character if number of read characters is higher than width 187 | // if it has met '\n' before maximum width has reached, it pads remaining bytes with space. 188 | func splitToLines(data []byte, width int) (lines [][]byte) { 189 | var pad int 190 | for i, w := 0, 0; i < len(data); i, w = i+1, w+1 { 191 | lineBreak := data[i] == '\n' 192 | switch { 193 | case lineBreak || w == width: 194 | source := data[pad:i] 195 | lines = append(lines, ensureWidth(source, width, ' ')) 196 | w = 0 197 | 198 | pad += len(source) 199 | if lineBreak { 200 | pad += 1 201 | } 202 | } 203 | } 204 | lines = append(lines, ensureWidth(data[pad:], width, ' ')) 205 | return 206 | } 207 | 208 | func ensureWidth(b []byte, w int, pad byte) []byte { 209 | if len(b) < w { 210 | s := make([]byte, len(b)) 211 | copy(s, b) 212 | return append(s, repeat(pad, w-len(b))...) 213 | } 214 | return b 215 | } 216 | 217 | func repeat(b byte, n int) []byte { 218 | ret := make([]byte, n) 219 | for i := 0; i < n; i++ { 220 | ret[i] = b 221 | } 222 | return ret 223 | } 224 | 225 | func max(a, b int) int { 226 | if a > b { 227 | return a 228 | } 229 | return b 230 | } 231 | -------------------------------------------------------------------------------- /lua/mod/stat/stat.go: -------------------------------------------------------------------------------- 1 | package stat 2 | 3 | import ( 4 | "fmt" 5 | "github.com/gobwas/gws/stat" 6 | "github.com/gobwas/gws/stat/counter/abs" 7 | "github.com/gobwas/gws/stat/counter/avg" 8 | "github.com/gobwas/gws/stat/counter/per" 9 | "github.com/yuin/gopher-lua" 10 | "time" 11 | ) 12 | 13 | type Stat struct { 14 | statistics *stat.Statistics 15 | } 16 | 17 | func New(s *stat.Statistics) *Stat { 18 | return &Stat{s} 19 | } 20 | 21 | func (m *Stat) Exports() lua.LGFunction { 22 | return func(L *lua.LState) int { 23 | mod := L.NewTable() 24 | mod.RawSetString("new", L.NewClosure(registerNew(m.statistics))) 25 | mod.RawSetString("abs", L.NewClosure(registerAbs(m.statistics))) 26 | mod.RawSetString("avg", L.NewClosure(registerAvg(m.statistics))) 27 | mod.RawSetString("per", L.NewClosure(registerPer(m.statistics))) 28 | mod.RawSetString("add", L.NewClosure(registerAdd(m.statistics))) 29 | mod.RawSetString("flush", L.NewClosure(registerFlush(m.statistics))) 30 | mod.RawSetString("pretty", L.NewClosure(registerPretty(m.statistics))) 31 | 32 | L.Push(mod) 33 | return 1 34 | } 35 | } 36 | 37 | const ( 38 | definitionFieldKind = "kind" 39 | definitionFieldMeta = "meta" 40 | definitionFieldInterval = "interval" 41 | ) 42 | 43 | const ( 44 | kindAbs int = iota 45 | kindAvg 46 | kindPer 47 | ) 48 | 49 | func registerNew(s *stat.Statistics) lua.LGFunction { 50 | return func(L *lua.LState) int { 51 | name := L.ToString(1) 52 | err := s.New(name) 53 | if err != nil { 54 | L.Push(lua.LString(err.Error())) 55 | return 1 56 | } 57 | 58 | for i := 2; ; i++ { 59 | def := L.ToTable(i) 60 | if def == nil { 61 | break 62 | } 63 | 64 | var meta map[string]interface{} 65 | m := def.RawGetString(definitionFieldMeta) 66 | if table, ok := m.(*lua.LTable); ok { 67 | meta = make(map[string]interface{}) 68 | table.ForEach(func(key lua.LValue, value lua.LValue) { 69 | if key.Type() == lua.LTString { 70 | meta[key.String()] = value 71 | } 72 | }) 73 | } 74 | 75 | var config stat.Config 76 | kind := int(def.RawGetString(definitionFieldKind).(lua.LNumber)) 77 | switch kind { 78 | case kindAbs: 79 | abs, err := absFactory() 80 | if err != nil { 81 | L.Push(lua.LString(err.Error())) 82 | return 1 83 | } 84 | config = stat.Config{abs, meta} 85 | 86 | case kindAvg: 87 | avg, err := avgFactory() 88 | if err != nil { 89 | L.Push(lua.LString(err.Error())) 90 | return 1 91 | } 92 | config = stat.Config{avg, meta} 93 | 94 | case kindPer: 95 | interval := def.RawGetString(definitionFieldInterval).(lua.LString) 96 | per, err := perFactory(interval.String()) 97 | if err != nil { 98 | L.Push(lua.LString(err.Error())) 99 | return 1 100 | } 101 | config = stat.Config{per, meta} 102 | 103 | default: 104 | L.Push(lua.LString(fmt.Sprintf("unknown type of counter: %v", kind))) 105 | return 1 106 | } 107 | 108 | err := s.Setup(name, config) 109 | if err != nil { 110 | L.Push(lua.LString(err.Error())) 111 | return 1 112 | } 113 | } 114 | 115 | return 0 116 | } 117 | } 118 | 119 | func registerPretty(s *stat.Statistics) lua.LGFunction { 120 | return func(L *lua.LState) int { 121 | L.Push(lua.LString(s.Pretty())) 122 | return 1 123 | } 124 | } 125 | 126 | func registerFlush(s *stat.Statistics) lua.LGFunction { 127 | return func(L *lua.LState) int { 128 | var index int 129 | results := L.NewTable() 130 | 131 | for name, metrics := range s.Metrics { 132 | for _, metric := range metrics { 133 | tags := L.NewTable() 134 | for key, value := range metric.Tags { 135 | tags.RawSetString(key, lua.LString(value)) 136 | } 137 | 138 | for _, instance := range metric.Instances { 139 | meta := L.NewTable() 140 | for key, value := range instance.Meta { 141 | meta.RawSetString(key, value.(lua.LValue)) 142 | } 143 | 144 | sub := L.NewTable() 145 | sub.RawSetString("name", lua.LString(name)) 146 | sub.RawSetString("kind", lua.LString(instance.Counter.Kind())) 147 | sub.RawSetString("value", lua.LNumber(instance.Counter.Flush())) 148 | sub.RawSetString("tags", tags) 149 | sub.RawSetString("meta", meta) 150 | 151 | results.RawSetInt(index, sub) 152 | index++ 153 | } 154 | } 155 | } 156 | L.Push(results) 157 | return 1 158 | } 159 | } 160 | 161 | func registerAdd(s *stat.Statistics) lua.LGFunction { 162 | return func(L *lua.LState) int { 163 | name := L.ToString(1) 164 | value := L.ToNumber(2) 165 | tags := L.Get(3) 166 | 167 | var tagsMap map[string]string 168 | if table, ok := tags.(*lua.LTable); ok { 169 | tagsMap = make(map[string]string) 170 | table.ForEach(func(key lua.LValue, value lua.LValue) { 171 | if key.Type() == lua.LTString { 172 | tagsMap[key.String()] = value.String() 173 | } 174 | }) 175 | } 176 | 177 | err := s.Increment(name, float64(value), tagsMap) 178 | if err != nil { 179 | L.Push(lua.LString(err.Error())) 180 | return 1 181 | } 182 | return 0 183 | } 184 | } 185 | func registerAbs(s *stat.Statistics) lua.LGFunction { 186 | return func(L *lua.LState) int { 187 | def := L.NewTable() 188 | def.RawSetString(definitionFieldKind, lua.LNumber(kindAbs)) 189 | if meta := L.ToTable(1); meta != nil { 190 | def.RawSetString(definitionFieldMeta, meta) 191 | } 192 | L.Push(def) 193 | return 1 194 | } 195 | } 196 | func registerAvg(s *stat.Statistics) lua.LGFunction { 197 | return func(L *lua.LState) int { 198 | def := L.NewTable() 199 | def.RawSetString(definitionFieldKind, lua.LNumber(kindAvg)) 200 | if meta := L.ToTable(1); meta != nil { 201 | def.RawSetString(definitionFieldMeta, meta) 202 | } 203 | L.Push(def) 204 | return 1 205 | } 206 | } 207 | func registerPer(s *stat.Statistics) lua.LGFunction { 208 | return func(L *lua.LState) int { 209 | def := L.NewTable() 210 | def.RawSetString(definitionFieldKind, lua.LNumber(kindPer)) 211 | def.RawSetString(definitionFieldInterval, L.Get(1)) 212 | if meta := L.ToTable(2); meta != nil { 213 | def.RawSetString(definitionFieldMeta, meta) 214 | } 215 | L.Push(def) 216 | return 1 217 | } 218 | } 219 | 220 | func absFactory() (stat.CounterFactory, error) { 221 | return func() stat.Counter { 222 | return abs.New() 223 | }, nil 224 | } 225 | func avgFactory() (stat.CounterFactory, error) { 226 | return func() stat.Counter { 227 | return avg.New() 228 | }, nil 229 | } 230 | func perFactory(dur string) (stat.CounterFactory, error) { 231 | duration, err := time.ParseDuration(dur) 232 | if err != nil { 233 | return nil, err 234 | } 235 | 236 | return func() stat.Counter { 237 | return per.New(duration) 238 | }, nil 239 | } 240 | -------------------------------------------------------------------------------- /server/server.go: -------------------------------------------------------------------------------- 1 | package server 2 | 3 | import ( 4 | "errors" 5 | "flag" 6 | "fmt" 7 | "github.com/chzyer/readline" 8 | "github.com/gobwas/gws/cli/color" 9 | "github.com/gobwas/gws/cli/input" 10 | "github.com/gobwas/gws/config" 11 | "github.com/gobwas/gws/ws" 12 | "github.com/gorilla/websocket" 13 | "io" 14 | "log" 15 | "net/http" 16 | "net/http/httputil" 17 | "os" 18 | "os/signal" 19 | "strconv" 20 | "strings" 21 | "sync" 22 | "sync/atomic" 23 | "time" 24 | ) 25 | 26 | var ( 27 | origin = flag.String("origin", "", "use this glob pattern for server origin checks") 28 | responder = &ResponderFlag{null, []string{echo, mirror, prompt, null}} 29 | ) 30 | 31 | func init() { 32 | flag.Var(responder, "response", fmt.Sprintf("how should server response on message (%s)", strings.Join(responder.expect, ", "))) 33 | } 34 | 35 | const ( 36 | echo = "echo" 37 | mirror = "mirror" 38 | prompt = "prompt" 39 | null = "null" 40 | ) 41 | 42 | func Go(c config.Config) error { 43 | var r Responder 44 | switch responder.Get() { 45 | case echo: 46 | r = EchoResponder 47 | case mirror: 48 | r = MirrorResponder 49 | case prompt: 50 | r = PromptResponder 51 | case null: 52 | r = DevNullResponder 53 | default: 54 | return errors.New("unknown responder type") 55 | } 56 | 57 | handler, err := newWsHandler(Config{ 58 | Headers: c.Headers, 59 | Origin: *origin, 60 | StatDump: c.StatDump, 61 | }, r) 62 | if err != nil { 63 | return err 64 | } 65 | 66 | handler.Init() 67 | 68 | log.Println("ready to listen", c.Addr) 69 | return http.ListenAndServe(c.Addr, handler) 70 | } 71 | 72 | type wsHandler struct { 73 | mu sync.Mutex 74 | 75 | upgrader ws.Upgrader 76 | config Config 77 | responder Responder 78 | sig chan os.Signal 79 | nextID uint64 80 | connsCount uint64 81 | conns map[uint64]connDescriptor 82 | 83 | requests uint64 84 | } 85 | 86 | type Config struct { 87 | Headers http.Header 88 | Origin string 89 | StatDump time.Duration 90 | } 91 | 92 | type connDescriptor struct { 93 | conn *websocket.Conn 94 | notice chan []byte 95 | done <-chan struct{} 96 | } 97 | 98 | const headerOrigin = "Origin" 99 | 100 | type Responder func(ws.Kind, []byte) ([]byte, error) 101 | 102 | func newWsHandler(c Config, r Responder) (*wsHandler, error) { 103 | return &wsHandler{ 104 | upgrader: ws.GetUpgrader(ws.UpgradeConfig{c.Origin, c.Headers}), 105 | config: c, 106 | responder: r, 107 | sig: make(chan os.Signal, 1), 108 | conns: make(map[uint64]connDescriptor), 109 | }, nil 110 | } 111 | 112 | func (h *wsHandler) Init() { 113 | signal.Notify(h.sig, os.Interrupt) 114 | go func() { 115 | listening: 116 | for _ = range h.sig { 117 | if h.connsCount == 0 { 118 | os.Exit(0) 119 | } 120 | 121 | h.mu.Lock() 122 | { 123 | var connId uint64 124 | if h.connsCount > 1 { 125 | var items []readline.PrefixCompleterInterface 126 | for id := range h.conns { 127 | items = append(items, readline.PcItem(string(id))) 128 | } 129 | completer := readline.NewPrefixCompleter(items...) 130 | 131 | r, err := input.ReadLine(&readline.Config{ 132 | Prompt: color.Green("> ") + "select connection id: ", 133 | AutoComplete: completer, 134 | }) 135 | if err != nil { 136 | if err == readline.ErrInterrupt { 137 | os.Exit(0) 138 | } 139 | log.Println("readline error:", err) 140 | continue listening 141 | } 142 | 143 | connId, err = strconv.ParseUint(string(r), 10, 64) 144 | if err != nil { 145 | log.Println("readline error:", err) 146 | continue listening 147 | } 148 | } else { 149 | connId = 1 150 | } 151 | 152 | r, err := input.ReadLine(&readline.Config{ 153 | Prompt: color.Green("> ") + fmt.Sprintf("notification for the connection #%d: ", connId), 154 | HistoryFile: "/tmp/gws_readline_server_notice.tmp", 155 | }) 156 | if err != nil { 157 | if err == readline.ErrInterrupt { 158 | os.Exit(0) 159 | } 160 | 161 | log.Println("readline error:", err) 162 | continue listening 163 | } 164 | 165 | h.conns[connId].notice <- r 166 | } 167 | h.mu.Unlock() 168 | } 169 | }() 170 | 171 | go func() { 172 | for range time.Tick(h.config.StatDump) { 173 | v := atomic.SwapUint64(&h.requests, 0) 174 | log.Printf("RPS: (%d) %.2f\n", v, float64(v/uint64(h.config.StatDump.Seconds()))) 175 | } 176 | }() 177 | } 178 | 179 | func (h *wsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { 180 | if config.Verbose { 181 | req, err := httputil.DumpRequest(r, false) 182 | if err != nil { 183 | log.Println(err) 184 | return 185 | } 186 | log.Println("new request", string(req)) 187 | } 188 | 189 | // take lock on read new connection 190 | h.mu.Lock() 191 | conn, err := h.upgrader(w, r) 192 | if err != nil { 193 | log.Println(err) 194 | h.mu.Unlock() 195 | return 196 | } 197 | h.connsCount++ 198 | h.nextID++ 199 | id := h.nextID 200 | desc := connDescriptor{ 201 | conn: conn, 202 | notice: make(chan []byte, 1), 203 | done: make(chan struct{}), 204 | } 205 | h.conns[id] = desc 206 | defer func() { 207 | conn.Close() 208 | h.mu.Lock() 209 | delete(h.conns, id) 210 | h.connsCount-- 211 | h.mu.Unlock() 212 | 213 | if config.Verbose { 214 | log.Printf("connection #%d closed\n", id) 215 | } 216 | }() 217 | h.mu.Unlock() 218 | 219 | if config.Verbose { 220 | log.Printf("establised connection #%d from %q\n", id, r.RemoteAddr) 221 | } 222 | 223 | in := ws.ReadAsyncFromConn(desc.done, conn) 224 | 225 | for { 226 | select { 227 | case <-desc.done: 228 | return 229 | 230 | case notice := <-desc.notice: 231 | err := ws.WriteToConn(conn, ws.TextMessage, notice) 232 | if err != nil { 233 | log.Println("error reading from socket:", err) 234 | return 235 | } 236 | log.Printf("sent message to %d: %s\n", id, string(notice)) 237 | 238 | case msg := <-in: 239 | atomic.AddUint64(&h.requests, 1) 240 | 241 | if msg.Err != nil { 242 | if msg.Err != io.EOF { 243 | log.Println("receive message error:", err) 244 | } 245 | 246 | return 247 | } 248 | if config.Verbose { 249 | log.Printf("received message from %d: %s\n", id, string(msg.Data)) 250 | } 251 | 252 | h.mu.Lock() 253 | resp, err := h.responder(msg.Kind, msg.Data) 254 | h.mu.Unlock() 255 | if err != nil { 256 | log.Println("responder error:", err) 257 | return 258 | } 259 | 260 | if resp != nil { 261 | err := ws.WriteToConn(conn, msg.Kind, resp) 262 | if err != nil { 263 | log.Println(err) 264 | return 265 | } 266 | 267 | if config.Verbose { 268 | log.Printf("sent message to %d: %s\n", id, string(resp)) 269 | } 270 | } 271 | } 272 | } 273 | } 274 | -------------------------------------------------------------------------------- /ev/ev.go: -------------------------------------------------------------------------------- 1 | package ev 2 | 3 | import ( 4 | "sync" 5 | "time" 6 | ) 7 | 8 | type RequestType int 9 | 10 | type Callback func(error, interface{}) 11 | 12 | type Handler interface { 13 | Init(*Loop) error 14 | Handle(*Loop, interface{}, Callback) error 15 | IsActive(*Loop) bool 16 | Stop() 17 | } 18 | 19 | type Loop struct { 20 | mu sync.Mutex 21 | 22 | handlers map[RequestType][]Handler 23 | requests []*request 24 | events []event 25 | teardowns []event 26 | timers []*Timer 27 | now time.Time 28 | done chan struct{} 29 | shutdown chan struct{} 30 | stop chan struct{} 31 | locked bool 32 | // idles []Idle 33 | } 34 | 35 | func NewLoop() *Loop { 36 | return &Loop{ 37 | done: make(chan struct{}), 38 | shutdown: make(chan struct{}), 39 | stop: make(chan struct{}, 1), 40 | handlers: make(map[RequestType][]Handler), 41 | now: time.Now(), 42 | } 43 | } 44 | 45 | func (l *Loop) Register(h Handler, t RequestType) { 46 | l.handlers[t] = append(l.handlers[t], h) 47 | if err := h.Init(l); err != nil { 48 | panic(err) 49 | } 50 | } 51 | 52 | func (l *Loop) Request(t RequestType, data interface{}, cb Callback) error { 53 | l.mu.Lock() 54 | { 55 | if !l.locked { 56 | l.requests = append(l.requests, &request{ 57 | t: t, 58 | cb: cb, 59 | data: data, 60 | }) 61 | } 62 | } 63 | l.mu.Unlock() 64 | return nil 65 | } 66 | 67 | func (l *Loop) Call(cb event) { 68 | l.mu.Lock() 69 | { 70 | if !l.locked { 71 | l.events = append(l.events, cb) 72 | } 73 | } 74 | l.mu.Unlock() 75 | } 76 | 77 | func (l *Loop) Timeout(delay time.Duration, repeat bool, cb event) *Timer { 78 | timer := &Timer{ 79 | delay: delay, 80 | repeat: repeat, 81 | cb: cb, 82 | } 83 | 84 | l.mu.Lock() 85 | { 86 | if !l.locked { 87 | timer.next = l.now.Add(timer.delay) 88 | l.timers = append(l.timers, timer) 89 | } 90 | } 91 | l.mu.Unlock() 92 | 93 | return timer 94 | } 95 | 96 | func (l *Loop) Teardown(cb event) { 97 | l.mu.Lock() 98 | { 99 | l.teardowns = append(l.teardowns, cb) 100 | } 101 | l.mu.Unlock() 102 | } 103 | 104 | func (l *Loop) Done() chan struct{} { 105 | return l.done 106 | } 107 | 108 | func (l *Loop) Shutdown() { 109 | close(l.shutdown) 110 | } 111 | 112 | func (l *Loop) Stop() { 113 | l.stop <- struct{}{} 114 | } 115 | 116 | func (l *Loop) lock() { 117 | l.mu.Lock() 118 | defer l.mu.Unlock() 119 | l.locked = true 120 | } 121 | 122 | func (l *Loop) Run() error { 123 | // l.mu.Lock() 124 | // defer l.mu.Unlock() 125 | // TODO(gobwas): what we should do if we run twice? 126 | 127 | go func(l *Loop) { 128 | for { 129 | select { 130 | case <-l.shutdown: 131 | close(l.done) 132 | return 133 | 134 | case <-l.stop: 135 | l.stopTimers() 136 | l.drainTeardown() 137 | l.stopHandlers() 138 | l.lock() 139 | 140 | default: 141 | if !l.IsAlive() { 142 | if !l.drainTeardown() { 143 | close(l.done) 144 | return 145 | } 146 | } else { 147 | l.updateNow() 148 | l.drainTimers() 149 | l.drainTicks() 150 | 151 | l.nextEvent() 152 | l.nextRequest() 153 | } 154 | } 155 | } 156 | }(l) 157 | 158 | return nil 159 | } 160 | 161 | func (l *Loop) IsAlive() bool { 162 | l.mu.Lock() 163 | defer l.mu.Unlock() 164 | 165 | if len(l.requests) > 0 { 166 | return true 167 | } 168 | if len(l.events) > 0 { 169 | return true 170 | } 171 | if len(l.timers) > 0 { 172 | return true 173 | } 174 | // if l.locked { 175 | // return false 176 | // } 177 | 178 | for _, handlers := range l.handlers { 179 | for _, handler := range handlers { 180 | if handler.IsActive(l) { 181 | return true 182 | } 183 | } 184 | } 185 | 186 | return false 187 | } 188 | 189 | func (l *Loop) updateNow() { 190 | l.now = time.Now() 191 | } 192 | 193 | func (l *Loop) stopHandlers() { 194 | l.mu.Lock() 195 | defer l.mu.Unlock() 196 | 197 | for _, handlers := range l.handlers { 198 | for _, h := range handlers { 199 | h.Stop() 200 | } 201 | } 202 | } 203 | 204 | func (l *Loop) stopTimers() { 205 | l.mu.Lock() 206 | defer l.mu.Unlock() 207 | 208 | for _, timer := range l.timers { 209 | timer.Stop() 210 | } 211 | } 212 | 213 | func (l *Loop) drainTimers() { 214 | var callbacks []event 215 | 216 | l.mu.Lock() 217 | { 218 | for i := 0; i < len(l.timers); { 219 | t := l.timers[i] 220 | 221 | var remove bool 222 | switch { 223 | case t.dropped: 224 | remove = true 225 | 226 | case t.next.Before(l.now): 227 | callbacks = append(callbacks, t.cb) 228 | 229 | if t.repeat { 230 | t.next = l.now.Add(t.delay) 231 | } else { 232 | remove = true 233 | } 234 | } 235 | 236 | if remove { 237 | l.deleteTimer(i) 238 | } else { 239 | i++ 240 | } 241 | } 242 | } 243 | l.mu.Unlock() 244 | 245 | for _, cb := range callbacks { 246 | cb() 247 | } 248 | } 249 | 250 | func (l *Loop) deleteTimer(i int) { 251 | copy(l.timers, l.timers[i+1:]) 252 | last := len(l.timers) - 1 253 | l.timers[last] = nil 254 | l.timers = l.timers[:last] 255 | } 256 | 257 | func (l *Loop) drainTicks() { 258 | // TODO(s.kamardin) 259 | } 260 | 261 | func (l *Loop) drainTeardown() bool { 262 | var callbacks []event 263 | 264 | l.mu.Lock() 265 | { 266 | callbacks = l.teardowns 267 | l.teardowns = nil 268 | } 269 | l.mu.Unlock() 270 | 271 | for _, cb := range callbacks { 272 | cb() 273 | } 274 | 275 | return len(callbacks) > 0 276 | } 277 | 278 | func (l *Loop) nextTeardown() bool { 279 | l.mu.Lock() 280 | 281 | if len(l.teardowns) == 0 { 282 | l.mu.Unlock() 283 | return false 284 | } 285 | 286 | teardown := l.teardowns[0] 287 | l.teardowns = append(l.teardowns[:0], l.teardowns[1:]...) 288 | 289 | l.mu.Unlock() 290 | 291 | teardown() 292 | return true 293 | } 294 | 295 | func (l *Loop) nextEvent() bool { 296 | l.mu.Lock() 297 | 298 | if len(l.events) == 0 { 299 | l.mu.Unlock() 300 | return false 301 | } 302 | 303 | evt := l.events[0] 304 | l.events = append(l.events[:0], l.events[1:]...) 305 | 306 | l.mu.Unlock() 307 | 308 | evt() 309 | 310 | return true 311 | } 312 | 313 | func (l *Loop) nextRequest() { 314 | l.mu.Lock() 315 | 316 | if len(l.requests) == 0 { 317 | l.mu.Unlock() 318 | return 319 | } 320 | 321 | evt := l.requests[0] 322 | l.requests = append(l.requests[:0], l.requests[1:]...) 323 | 324 | l.mu.Unlock() 325 | 326 | { 327 | for _, handler := range l.handlers[evt.t] { 328 | err := handler.Handle(l, evt.data, evt.cb) 329 | if err != nil { 330 | panic(err) 331 | } 332 | } 333 | } 334 | } 335 | 336 | type event func() 337 | 338 | type request struct { 339 | t RequestType 340 | cb Callback 341 | data interface{} 342 | } 343 | 344 | func (e *request) Data() interface{} { 345 | return e.data 346 | } 347 | 348 | func (e *request) Callback() Callback { 349 | return e.cb 350 | } 351 | 352 | type Timer struct { 353 | delay time.Duration 354 | repeat bool 355 | cb event 356 | dropped bool 357 | next time.Time 358 | } 359 | 360 | func (t *Timer) Stop() { 361 | t.dropped = true 362 | } 363 | --------------------------------------------------------------------------------