chore: add SetupContextForConn for common/net

This commit is contained in:
wwqgtxx 2023-11-03 11:00:40 +08:00
parent 885ee7a820
commit ee3038d5e4

31
common/net/context.go Normal file
View file

@ -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
}
}
}