From 80393295c1185e50d0b784d4bc5ffaa918d187b9 Mon Sep 17 00:00:00 2001
From: LiuYang <kfliuyang@gmail.com>
Date: Thu, 17 Aug 2023 23:25:09 +0800
Subject: [PATCH] Correct way to save memory using write buffer pool and
 freeing net.http default buffers (#761)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

**Summary of Changes**

1. Add an example that uses the write buffer pool

The loop process of the websocket connection is inner the http handler
at existing examples， This usage will cause the 8k buffer(4k read buffer
+ 4k write buffer) allocated by net.http can't be GC(Observed by heap
profiling, see picture below) . The purpose of saving memory is not
achieved even if the WriteBufferPool is used.

In example bufferpool, server process websocket connection in a new
goroutine, and the goroutine created by the net.http will exit, then the
8k buffer will be GC.


![heap](https://user-images.githubusercontent.com/12793501/148676918-872d1a6d-ce10-4146-ba01-7de114db09f5.png)

Co-authored-by: hakunaliu <hakunaliu@tencent.com>
Co-authored-by: Corey Daley <cdaley@redhat.com>
---
 README.md                     |  1 +
 examples/bufferpool/client.go | 89 +++++++++++++++++++++++++++++++++++
 examples/bufferpool/server.go | 55 ++++++++++++++++++++++
 3 files changed, 145 insertions(+)
 create mode 100644 examples/bufferpool/client.go
 create mode 100644 examples/bufferpool/server.go

diff --git a/README.md b/README.md
index d33ed7fd..9b6956f3 100644
--- a/README.md
+++ b/README.md
@@ -14,6 +14,7 @@ Gorilla WebSocket is a [Go](http://golang.org/) implementation of the
 * [Command example](https://github.com/gorilla/websocket/tree/master/examples/command)
 * [Client and server example](https://github.com/gorilla/websocket/tree/master/examples/echo)
 * [File watch example](https://github.com/gorilla/websocket/tree/master/examples/filewatch)
+* [Write buffer pool example](https://github.com/gorilla/websocket/tree/master/examples/bufferpool)
 
 ### Status
 
diff --git a/examples/bufferpool/client.go b/examples/bufferpool/client.go
new file mode 100644
index 00000000..a3719a9a
--- /dev/null
+++ b/examples/bufferpool/client.go
@@ -0,0 +1,89 @@
+//go:build ignore
+// +build ignore
+
+package main
+
+import (
+	"flag"
+	"log"
+	"net/url"
+	"os"
+	"os/signal"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+var addr = flag.String("addr", "localhost:8080", "http service address")
+
+func runNewConn(wg *sync.WaitGroup) {
+	defer wg.Done()
+
+	interrupt := make(chan os.Signal, 1)
+	signal.Notify(interrupt, os.Interrupt)
+
+	u := url.URL{Scheme: "ws", Host: *addr, Path: "/ws"}
+	log.Printf("connecting to %s", u.String())
+	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+	if err != nil {
+		log.Fatal("dial:", err)
+	}
+	defer c.Close()
+
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+		for {
+			_, message, err := c.ReadMessage()
+			if err != nil {
+				log.Println("read:", err)
+				return
+			}
+			log.Printf("recv: %s", message)
+		}
+	}()
+
+	ticker := time.NewTicker(time.Minute * 5)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-done:
+			return
+		case t := <-ticker.C:
+			err := c.WriteMessage(websocket.TextMessage, []byte(t.String()))
+			if err != nil {
+				log.Println("write:", err)
+				return
+			}
+		case <-interrupt:
+			log.Println("interrupt")
+
+			// Cleanly close the connection by sending a close message and then
+			// waiting (with timeout) for the server to close the connection.
+			err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+			if err != nil {
+				log.Println("write close:", err)
+				return
+			}
+			select {
+			case <-done:
+			case <-time.After(time.Second):
+			}
+			return
+		}
+	}
+}
+
+func main() {
+	flag.Parse()
+	log.SetFlags(0)
+	wg := &sync.WaitGroup{}
+	for i := 0; i < 1000; i++ {
+		wg.Add(1)
+		go runNewConn(wg)
+	}
+	wg.Wait()
+}
diff --git a/examples/bufferpool/server.go b/examples/bufferpool/server.go
new file mode 100644
index 00000000..25bb20f5
--- /dev/null
+++ b/examples/bufferpool/server.go
@@ -0,0 +1,55 @@
+//go:build ignore
+// +build ignore
+
+package main
+
+import (
+	"flag"
+	"log"
+	"net/http"
+	"sync"
+
+	_ "net/http/pprof"
+
+	"github.com/gorilla/websocket"
+)
+
+var addr = flag.String("addr", "localhost:8080", "http service address")
+
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  256,
+	WriteBufferSize: 256,
+	WriteBufferPool: &sync.Pool{},
+}
+
+func process(c *websocket.Conn) {
+	defer c.Close()
+	for {
+		_, message, err := c.ReadMessage()
+		if err != nil {
+			log.Println("read:", err)
+			break
+		}
+		log.Printf("recv: %s", message)
+	}
+}
+
+func handler(w http.ResponseWriter, r *http.Request) {
+	c, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Print("upgrade:", err)
+		return
+	}
+
+	// Process connection in a new goroutine
+	go process(c)
+
+	// Let the http handler return, the 8k buffer created by it will be garbage collected
+}
+
+func main() {
+	flag.Parse()
+	log.SetFlags(0)
+	http.HandleFunc("/ws", handler)
+	log.Fatal(http.ListenAndServe(*addr, nil))
+}