From ee3038d5e4e80fcc82cd849c30a5f6801af4607d Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Fri, 3 Nov 2023 11:00:40 +0800
Subject: [PATCH] chore: add SetupContextForConn for common/net

---
 common/net/context.go | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)
 create mode 100644 common/net/context.go

diff --git a/common/net/context.go b/common/net/context.go
new file mode 100644
index 000000000..917028d16
--- /dev/null
+++ b/common/net/context.go
@@ -0,0 +1,31 @@
+package net
+
+import (
+	"context"
+	"net"
+)
+
+// SetupContextForConn is a helper function that starts connection I/O interrupter goroutine.
+func SetupContextForConn(ctx context.Context, conn net.Conn) (done func(*error)) {
+	var (
+		quit      = make(chan struct{})
+		interrupt = make(chan error, 1)
+	)
+	go func() {
+		select {
+		case <-quit:
+			interrupt <- nil
+		case <-ctx.Done():
+			// Close the connection, discarding the error
+			_ = conn.Close()
+			interrupt <- ctx.Err()
+		}
+	}()
+	return func(inputErr *error) {
+		close(quit)
+		if ctxErr := <-interrupt; ctxErr != nil && inputErr != nil {
+			// Return context error to user.
+			inputErr = &ctxErr
+		}
+	}
+}