wsjson: Write messages in a single frame always

Closes #315

Author: nhooyr

internal/util/util.go

@@ -0,0 +1,7 @@
+package util
+
+type WriterFunc func(p []byte) (int, error)
+
+func (f WriterFunc) Write(p []byte) (int, error) {
+	return f(p)
+}

wsjson/wsjson.go

@@ -8,6 +8,7 @@ import (
 
 	"nhooyr.io/websocket"
 	"nhooyr.io/websocket/internal/bpool"
+	"nhooyr.io/websocket/internal/util"
 	"nhooyr.io/websocket/internal/errd"
 )
 
@@ -51,17 +52,17 @@ func Write(ctx context.Context, c *websocket.Conn, v interface{}) error {
 func write(ctx context.Context, c *websocket.Conn, v interface{}) (err error) {
 	defer errd.Wrap(&err, "failed to write JSON message")
 
-	w, err := c.Writer(ctx, websocket.MessageText)
-	if err != nil {
-		return err
-	}
-
 	// json.Marshal cannot reuse buffers between calls as it has to return
 	// a copy of the byte slice but Encoder does as it directly writes to w.
-	err = json.NewEncoder(w).Encode(v)
+	err = json.NewEncoder(util.WriterFunc(func(p []byte) (int, error) {
+		err := c.Write(ctx, websocket.MessageText, p)
+		if err != nil {
+			return 0, err
+		}
+		return len(p), nil
+	})).Encode(v)
 	if err != nil {
 		return fmt.Errorf("failed to marshal JSON: %w", err)
 	}
-
-	return w.Close()
+	return nil
 }