Write a simple WebSocket push service with Go

The code for this article can be found at github.com/alfred-zhon… To obtain.

background

I recently got a request to display alarm information on the web page. In the past, alarm information was pushed to users through SMS, wechat and App. Now, login users can also receive alarm push in real time on the web.

I vaguely remember having a similar need in my previous job. Comet4J was used to solve this problem because the browser standards were older and Java was used more at the time. The specific principle is long polling, long links. But now that HTML5 is popular and IE has been replaced by Edge, the technology seems outdated.

I’ve heard of WebSocket for a long time, but since many users didn’t have browser support at the time, I’ve only dabbled with the technology without studying it too deeply. Now I’m going to take a look at it while the project needs it.

Introduction of websocket

In the past, browsers used to get data from the server by sending HTTP requests and waiting for the server to respond. That is to say, the browser side is always the initiator of the whole request, and only it takes the initiative to get the data. In order for the browser side to get real-time data from the server side, it needs to make continuous requests to the server side. Although most of the time no actual data was retrieved, this increased the network strain significantly, and the strain on the server side skyrocketed.

Later we learned to use long join + long poll. In other words, keep HTTP requests alive for as long as possible and keep HTTP connections alive. This takes a lot of stress out of it, but it still requires constant polling and is not really real-time. (Borrow a picture)

WebSocket was made standard in 2011 with the advent of HTML5 (see RFC 6455 for details).

To borrow from Go Web Programming. WebSocket uses special headers that allow the browser and the server to establish a connection between the browser and the server by simply shaking hands. While the connection remains active, you can use JavaScript to write to or receive data from the connection, just like using a regular TCP Socket. It solves the problem of making the Web real-time.

Since WebSocket is full-duplex communication, after a WebSocket connection is established, the subsequent communication is similar to traditional TCP communication. Clients and servers can send data to each other without any real time problems.

Development package selection

WebSocket support is not included in the official Go SDK, so third-party libraries must be used.

To develop websockets using Golang, the choice is basically between X /net/ WebSocket and Gorilla/WebSocket. The examples in The book Go Web Programming use X /net/websocket as a development package, and it seems more official and formal. Not so, according to the feedback I’ve received from online searches. X /net/websocket seems to be buggy and unstable, and the problem is not solved in a timely manner. Gorilla/Websocket is superior in comparison.

And contributions to the Gorilla Web Toolkit organization must be appreciated. 🙏. There are not only WebSocket implementations, but also some other tools. Everyone is welcome to use it and feel free to give feedback or contributions.

Implementation of Push Service

The basic principle of

The preliminary design of the project is as follows:

After the server starts, two handlers are registered.

• The websocketHandler is used to enable the browser to send an Upgrade request and Upgrade to a WebSocket connection.
• PushHandler is used to provide a request for an external push to send push data.

The browser first connects to the websocketHandler (default address is WS :// IP :port/ws). The upgrade request is a WebSocket connection. After the connection is established, the user needs to send the registration information for registration. The registration information contains a token information. The server validates the supplied token and obtains the corresponding userId (generally, a userId may be associated with many tokens), and maintains the relationship between the token, userId, and CONN (connection).

PushHandler (default address ws:// IP :port/push) sends a request to pushHandler that contains the userId and message fields. The server retrieves all conns connected to the server based on the userId and pushes messages one by one.

Due to the real-time nature of the push service, the data pushed does not and does not need to be cached.

Code,

I’ll talk a little bit about the basics of the code here, as well as some common writing and patterns in Go. (I’ve moved to Go from other languages, and Go is pretty young, too. So if you have any suggestions, feel free to do so.) . Because most of the inventors and major maintainers of Go come from C/C++, the code of Go tends to be in the C/C++ family.

Let’s look at the Server structure first:

// Server defines parameters for running websocket server.
type Server struct {
	// Address for server to listen on
	Addr string

	// Path for websocket request, default "/ws".
	WSPath string

	// Path for push message, default "/push".
	PushPath string

	// Upgrader is for upgrade connection to websocket connection using
	// "github.com/gorilla/websocket".
	//
	// If Upgrader is nil, default upgrader will be used. Default upgrader is
	// set ReadBufferSize and WriteBufferSize to 1024, and CheckOrigin always
	// returns true.
	Upgrader *websocket.Upgrader

	// Check token if it's valid and return userID. If token is valid, userID
	// must be returned and ok should be true. Otherwise ok should be false.
	AuthToken func(token string) (userID string, ok bool)

	// Authorize push request. Message will be sent if it returns true, / /otherwise the request will be discarded. Default nil and push request
	// will always be accepted.
	PushAuth func(r *http.Request) bool

	wh *websocketHandler
	ph *pushHandler
}
Copy the code

PS: As my entire project notes are written in English, so forgive me, I hope it does not interfere with reading.

Upgrader * webSocket. Upgrader, the object of the Gorilla/Websocket package, is used to upgrade HTTP requests.

If a structure has too many parameters, it is usually not recommended to initialize it directly, but to use the New method it provides. Here is:

// NewServer creates a new Server.
func NewServer(addr string) *Server {
	return &Server{
		Addr:     addr,
		WSPath:   serverDefaultWSPath,
		PushPath: serverDefaultPushPath,
	}
}
Copy the code

This is also a common use of the Go language to provide external initialization methods.

The Server then starts using the ListenAndServe method and listens on the port, similar to the use of HTTP packets:

// ListenAndServe listens on the TCP network address and handle websocket
// request.
func (s *Server) ListenAndServe(a) error {
	b := &binder{
		userID2EventConnMap: make(map[string]*[]eventConn),
		connID2UserIDMap:    make(map[string]string),}// websocket request handler
	wh := websocketHandler{
		upgrader: defaultUpgrader,
		binder:   b,
	}
	ifs.Upgrader ! =nil {
		wh.upgrader = s.Upgrader
	}
	ifs.AuthToken ! =nil {
		wh.calcUserIDFunc = s.AuthToken
	}
	s.wh = &wh
	http.Handle(s.WSPath, s.wh)

	// push request handler
	ph := pushHandler{
		binder: b,
	}
	ifs.PushAuth ! =nil {
		ph.authFunc = s.PushAuth
	}
	s.ph = &ph
	http.Handle(s.PushPath, s.ph)

	return http.ListenAndServe(s.Addr, nil)}Copy the code

Here we generate two handlers, websocketHandler and pushHandler. The websocketHandler connects to the browser and transmits data, while the pushHandler handles the push requests. As you can see, both handlers encapsulates a Binder object. This binder is used to maintain the relationship between token <-> userID <-> Conn:

// binder is defined to store the relation of userID and eventConn
type binder struct {
	mu sync.RWMutex

	// map stores key: userID and value of related slice of eventConn
	userID2EventConnMap map[string]*[]eventConn

	// map stores key: connID and value: userID
	connID2UserIDMap map[string]string
}
Copy the code

websocketHandler

Take a look at the implementation of websocketHandler.

// websocketHandler defines to handle websocket upgrade request.
type websocketHandler struct {
	// upgrader is used to upgrade request.
	upgrader *websocket.Upgrader

	// binder stores relations about websocket connection and userID.
	binder *binder

	// calcUserIDFunc defines to calculate userID by token. The userID will
	// be equal to token if this function is nil.
	calcUserIDFunc func(token string) (userID string, ok bool)
}
Copy the code

Very simple structure. The websocketHandler implements the HTTP.handler interface:

// First try to upgrade connection to websocket. If success, connection will
// be kept until client send close message or server drop them.
func (wh *websocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	wsConn, err := wh.upgrader.Upgrade(w, r, nil)
	iferr ! =nil {
		return
	}
	defer wsConn.Close()

	// handle Websocket request
	conn := NewConn(wsConn)
	conn.AfterReadFunc = func(messageType int, r io.Reader) {
		var rm RegisterMessage
		decoder := json.NewDecoder(r)
		iferr := decoder.Decode(&rm); err ! =nil {
			return
		}

		// calculate userID by token
		userID := rm.Token
		ifwh.calcUserIDFunc ! =nil {
			uID, ok := wh.calcUserIDFunc(rm.Token)
			if! ok {return
			}
			userID = uID
		}

		// bind
		wh.binder.Bind(userID, rm.Event, conn)
	}
	conn.BeforeCloseFunc = func(a) {
		// unbind
		wh.binder.Unbind(conn)
	}

	conn.Listen()
}
Copy the code

We first convert the incoming HTTP. Request to websocket.Conn and then unpack it into our own wServer. Conn (wrapper, or combination, is a typical use of the Go language. Remember, Go has no inheritance, only composition). Conn’s AfterReadFunc and BeforeCloseFunc methods are then set and conn.listen () is started. AfterReadFunc means that when Conn reads the data, it tries to validate and calculate the userID based on the token, and then registers the binding with bind. BeforeCloseFunc unforeclosefunc is foreclosed before Conn closes.

pushHandler

PushHandler is easier to understand. It parses the request and pushes data:

// Authorize if needed. Then decode the request and push message to each
// realted websocket connection.
func (s *pushHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ifr.Method ! = http.MethodPost { w.WriteHeader(http.StatusMethodNotAllowed)return
	}

	// authorize
	ifs.authFunc ! =nil {
		ifok := s.authFunc(r); ! ok { w.WriteHeader(http.StatusUnauthorized)return}}// read request
	var pm PushMessage
	decoder := json.NewDecoder(r.Body)
	iferr := decoder.Decode(&pm); err ! =nil {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(ErrRequestIllegal.Error()))
		return
	}

	// validate the data
	if pm.UserID == "" || pm.Event == "" || pm.Message == "" {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(ErrRequestIllegal.Error()))
		return
	}

	cnt, err := s.push(pm.UserID, pm.Event, pm.Message)
	iferr ! =nil {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(err.Error()))
		return
	}

	result := strings.NewReader(fmt.Sprintf("message sent to %d clients", cnt))
	io.Copy(w, result)
}
Copy the code

Conn

Conn (wserver.Conn) is a wrapper for websocket.Conn.

// Conn wraps websocket.Conn with Conn. It defines to listen and read
// data from Conn.
type Conn struct {
	Conn *websocket.Conn

	AfterReadFunc   func(messageType int, r io.Reader)
	BeforeCloseFunc func(a)

	once   sync.Once
	id     string
	stopCh chan struct{}}Copy the code

The main method is Listen() :

// Listen listens for receive data from websocket connection. It blocks
// until websocket connection is closed.
func (c *Conn) Listen(a) {
	c.Conn.SetCloseHandler(func(code int, text string) error {
		ifc.BeforeCloseFunc ! =nil {
			c.BeforeCloseFunc()
		}

		iferr := c.Close(); err ! =nil {
			log.Println(err)
		}

		message := websocket.FormatCloseMessage(code, "")
		c.Conn.WriteControl(websocket.CloseMessage, message, time.Now().Add(time.Second))
		return nil
	})

	// Keeps reading from Conn util get error.
ReadLoop:
	for {
		select {
		case <-c.stopCh:
			break ReadLoop
		default:
			messageType, r, err := c.Conn.NextReader()
			iferr ! =nil {
				// TODO: handle read error maybe
				break ReadLoop
			}

			ifc.AfterReadFunc ! =nil {
				c.AfterReadFunc(messageType, r)
			}
		}
	}
}
Copy the code

The main setup is to handle and keep reading data when the WebSocket connection is closed.

It is difficult to describe the operation process of the entire code in a comprehensive way. For detailed reading of the code, please visit github.com/alfred-zhon… To obtain.

Afterword.

The code has been tested and has been running in a formal environment for a while. But the code may still be unstable, so it’s not unusual for problems to occur during use. Feel free to send me issues or PRs.

reference

• Go Web Programming — Astaxie
• Long polling (HOOJO) for Web communication
• Gorilla web toolkit