A Million WebSocket and Go

1
2
3
4
5

+--------------+     (2)    +-------------+      (1)    +-----------+
|              | <--------+ |             |  <--------+ |           |
|    Storage   |            |     API     |     HTTP    |  Browser  |
|              | +--------> |             |  +--------> |           |
+--------------+     (3)    +-------------+      (4)    +-----------+

1
2
3
4
5
6
7
8
9

+--------------+            +-------------+   WebSocket  +-----------+
|    Storage   |            |     API     | +----------> |  Browser  |
+--------------+            +-------------+      (3)     +-----------+
       +                           ^
       | (1)                       | (2)
       v                           +
+-----------------------------------------+
|                  Bus                    |
+-----------------------------------------+

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

// Packet represents application level data.
type Packet struct {
    
}

// Channel wraps user connection.
type Channel struct {
    conn net.Conn    // WebSocket connection
    send chan Packet // Outgoing packets queue
}

func NewChannel(conn net.Conn) *Channel {
    c := &Channel{
        conn: conn,
        send: make(chan Packet, N),
    }
    go c.reader()
    go c.writer()
	return c
}

1
2
3
4
5
6
7
8

func (c *Channel) reader() {
    // We make a buffered read to reduce read syscalls.
    buf := bufio.NewReader(c.conn)
    for {
        pkt, _ := readPacket(buf)
        c.handle(pkt)
    }
}

1
2
3
4
5
6
7
8
9

func (c *Channel) writer() {
    // we make buffered write to reduce write syscalls.
    buf := bufio.NewWriter(c.conn)
    
    for pkt := range c.send {
        _ := writePacket(buf, pkt)
        buf.Flush()
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

import (
    "net/http"
    "some/websocket"
)

http.HandleFunc("/v1/ws", func(w http.ResponseWriter, r *http.Request) {
    conn, _ := websocket.Upgrade(r, w)
    ch := NewChannel(conn)
    // ...
})

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

// net/fd_unix.go

func (fd *netFD) Read(p []byte) (n int, err error) {
    // ...
    for {
        n, err = syscall.Read(fd.sysfd, p)
        if err != nil {
            n = 0
            if err == syscall.EAGAIN {
                if err = fd.pd.waitRead(); err == nil {
                    continue
                }
            }
        }
        // ...
        break
    }
    // ...
}

1
2
3
4
5
6
7
8

// net/fd_poll_runtime.go

func (pd *pollDesc) waitRead() error {
    return pd.wait('r')
}
func (pd *pollDesc) wait(mode int) error {
    res := runtime_pollWait(pd.runtimeCtx, mode)
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

ch := NewChannel(conn)

// Make conn to be observed by netpoll instance.
poller.Start(conn, netpoll.EventRead, func() {
    // We spawn goroutine here to prevent poller wait loop
    // to become locked during receiving packet from ch
    go Receive(ch)
})

func (ch *Channel) Receive() {
    buf := bufio.NewReader(ch.conn)
    pkt := readPacket(buf)
    c.handle(pkt)
}

1
2
3
4
5
6

func (ch *Channel) Send(p Packet) {
    if c.noWriterYet() {
        go ch.writer()
    }
    ch.send <- p
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

package gopool

func New(size int) *Pool {
    reutrn &Pool{
        work: make(chan func()),
        sem:  make(chan struct{}, size),
    }
}

func (p *Pool) Schedule(task func()) error {
    select {
    case p.work <- task:
    case p.sem <- struct{}{}:
        go p.worker(task)    
    }
}

func (p *Pool) worker(task func()) {
    defer func() { <- p.sem }
    for {
        task()
        task = <- p.work
    }
}

1
2
3
4
5
6
7

pool := gopool.New(128)
pooler.Start(conn, netpool.EventRead, func() {
    // We will block pooler wait loop when all pool workers are busy
    pool.Schedule(func() {
        Receive(ch)
    })
})

1
2
3
4
5
6
7

pool := gopool.New(128)
func (ch *Channel) Send(p Packet) {
    if c.noWriterYet() {
        pool.Schedule(ch.writer)
    }
    ch.send <- p
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

GET /ws HTTP/1.1
Host: Mail.Ru
Connection: Upgrade
Sec-Websocket-Key: A3xNe7sEB9HixkmBhVrYaA==
Sec-Websocket-Version: 13
Upgrade: websocket

HTTP/1.1 101 Switching Protocols
Connection: Upgrade
Sec-Websocket-Accept: ksu0wXWG+YmkVx+KQR2agP0cQn4=
Upgrade: websocket

1
2

func ReadFrame(ioReader) (Frame, error)
func WriteFrame(io.Writer, Frame) error

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

// getReadBuf, putReadBuf are intended to 
// reuse *bufio.Reader (with sync.Pool for example).
func getReadBuf(io.Reader) *bufio.Reader
func putReadBuf(*bufio.Reader)

// readPacket must be called when data could be read from conn.
func readPacket(conn io.Reader) error {
    buf := getReadBuf()
    defer putReadBuf(buf)

    buf.Reset(conn)
    frame, _ := ReadFrame(buf)
    parsePacket(frame.Payload)
    //...
}

1
2

BenchmarkUpgradeHTTP    5156 ns/op    8576 B/op    9 allocs/op
BenchmarkUpgradeTCP     973 ns/op     0 B/op       0 allocs/op

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

import (
    "net"
    "github.com/gobwas/ws"
)

ln, _ := net.Listen("tcp", ":8080")

for {
    // Try to accept incoming connection inside free pool worker.
    // If there no free workers for 1ms, do not accept anything and try later.
    // This will help us to prevent many self-ddos or out of resource limit cases.
    err := pool.ScheduleTimeout(time.Millisecond, func() {
        conn := ln.Accept()
        _ = ws.Upgrade(conn)

        // Wrap WebSocket connection with our Channel struct.
        // This will help us to handle/send our app's packets.
        ch := NewChannel(conn)

        // Wait for incoming bytes from connection.
        poller.Start(conn, netpoll.EventRead, func() {
            // Do not cross the resource limits.
            pool.Schedule(func() {
                // Read and handle incoming packet(s).
                ch.Recevie()
            })
        })
    })
    if err != nil {   
        time.Sleep(time.Millisecond)
    }
}