event/event_test: add BenchmarkReceive

event/event.go

@@ -39,7 +39,10 @@ func NewClient(socket string) (*EventClient, error) {
 	if err != nil {
 		return nil, fmt.Errorf("error while connecting to socket: %w", err)
 	}
-	return &EventClient{conn: conn}, err
+	return &EventClient{
+		conn:   conn,
+		reader: bufio.NewReaderSize(conn, bufSize),
+	}, err
 }
 
 // Close the underlying connection.
@@ -56,7 +59,7 @@ func (c *EventClient) Close() error {
 func (c *EventClient) Receive(ctx context.Context) ([]ReceivedData, error) {
 	buf := make([]byte, bufSize)
 
-	n, err := readWithContext(ctx, c.conn, buf)
+	n, err := c.readWithContext(ctx, buf)
 	if err != nil {
 		return nil, fmt.Errorf("error while reading from socket: %w", err)
 	}
@@ -96,15 +99,14 @@ func (c *EventClient) Subscribe(ctx context.Context, ev EventHandler, events ...
 	}
 }
 
-func readWithContext(ctx context.Context, conn net.Conn, buf []byte) (int, error) {
+func (c *EventClient) readWithContext(ctx context.Context, buf []byte) (int, error) {
 	done := make(chan struct{})
 	var n int
 	var err error
 
 	// Start a goroutine to perform the read
 	go func() {
-		reader := bufio.NewReader(conn)
-		n, err = reader.Read(buf)
+		n, err = c.reader.Read(buf)
 		close(done)
 	}()
 
@@ -113,9 +115,9 @@ func readWithContext(ctx context.Context, conn net.Conn, buf []byte) (int, error
 		return n, err
 	case <-ctx.Done():
 		// Set a short deadline to unblock the Read()
-		conn.SetReadDeadline(time.Now())
+		c.conn.SetReadDeadline(time.Now())
 		// Reset read deadline
-		defer conn.SetReadDeadline(time.Time{})
+		defer c.conn.SetReadDeadline(time.Time{})
 		// Make sure that the goroutine is done to avoid leaks
 		<-done
 		return 0, ctx.Err()

event/event_test.go

@@ -1,8 +1,11 @@
 package event
 
 import (
+	"bufio"
 	"context"
 	"errors"
+	"math/rand"
+	"net"
 	"os"
 	"testing"
 	"time"
@@ -11,6 +14,8 @@ import (
 	"github.com/thiagokokada/hyprland-go/internal/assert"
 )
 
+const socketPath = "/tmp/bench_unix_socket.sock"
+
 type FakeEventClient struct {
 	EventClient
 }
@@ -245,3 +250,91 @@ func (h *FakeEventHandler) Screencast(s Screencast) {
 	assert.Equal(h.t, s.Owner, "0")
 	assert.Equal(h.t, s.Sharing, true)
 }
+
+func BenchmarkReceive(b *testing.B) {
+	go RandomStringServer()
+
+	// Make sure the socket exist
+	for i := 0; i < 10; i++ {
+		time.Sleep(100 * time.Millisecond)
+		if _, err := os.Stat(socketPath); err != nil {
+			break
+		}
+	}
+
+	c := assert.Must1(NewClient(socketPath))
+	defer c.Close()
+
+	ctx := context.Background()
+
+	// Reset setup time
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		c.Receive(ctx)
+	}
+}
+
+// This function needs to be as fast as possible, otherwise this is the
+// bottleneck
+// https://stackoverflow.com/a/31832326
+const (
+	letterBytes   = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+	letterIdxBits = 6                    // 6 bits to represent a letter index
+	letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
+	letterIdxMax  = 63 / letterIdxBits   // # of letter indices fitting in 63 bits
+)
+
+func RandomBytes(n int) []byte {
+	b := make([]byte, n)
+	// A rand.Int63() generates 63 random bits, enough for letterIdxMax letters!
+	for i, cache, remain := n-1, rand.Int63(), letterIdxMax; i >= 0; {
+		if remain == 0 {
+			cache, remain = rand.Int63(), letterIdxMax
+		}
+		if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
+			b[i] = letterBytes[idx]
+			i--
+		}
+		cache >>= letterIdxBits
+		remain--
+	}
+
+	return b
+}
+
+func RandomStringServer() {
+	// Remove the previous socket file if it exists
+	if err := os.RemoveAll(socketPath); err != nil {
+		panic(err)
+	}
+
+	listener, err := net.Listen("unix", socketPath)
+	if err != nil {
+		panic(err)
+	}
+	defer listener.Close()
+
+	for {
+		conn, err := listener.Accept()
+		if err != nil {
+			panic(err)
+		}
+		writer := bufio.NewWriter(conn)
+
+		go func(c net.Conn) {
+			defer c.Close()
+
+			for {
+				prefix := []byte(">>>")
+				randomData := RandomBytes(16)
+				message := append(prefix, randomData...)
+
+				// Send the message to the client
+				_, err := writer.Write(message)
+				if err != nil {
+					return
+				}
+			}
+		}(conn)
+	}
+}

event/event_types.go

@@ -1,13 +1,15 @@
 package event
 
 import (
+	"bufio"
 	"context"
 	"net"
 )
 
 // EventClient is the event struct from hyprland-go.
 type EventClient struct {
-	conn net.Conn
+	conn   net.Conn
+	reader *bufio.Reader
 }
 
 // Event Client interface, right now only used for testing.