package connctx
import (
"context"
"errors"
"io"
"net"
"sync"
"sync/atomic"
"time"
)
var ErrClosing = errors .New ("use of closed network connection" )
type Reader interface {
ReadContext (context .Context , []byte ) (int , error )
}
type Writer interface {
WriteContext (context .Context , []byte ) (int , error )
}
type ReadWriter interface {
Reader
Writer
}
type ConnCtx interface {
Reader
Writer
io .Closer
LocalAddr () net .Addr
RemoteAddr () net .Addr
Conn () net .Conn
}
type connCtx struct {
nextConn net .Conn
closed chan struct {}
closeOnce sync .Once
readMu sync .Mutex
writeMu sync .Mutex
}
var veryOld = time .Unix (0 , 1 )
func New (conn net .Conn ) ConnCtx {
c := &connCtx {
nextConn : conn ,
closed : make (chan struct {}),
}
return c
}
func (c *connCtx ) ReadContext (ctx context .Context , b []byte ) (int , error ) {
c .readMu .Lock ()
defer c .readMu .Unlock ()
select {
case <- c .closed :
return 0 , io .EOF
default :
}
done := make (chan struct {})
var wg sync .WaitGroup
var errSetDeadline atomic .Value
wg .Add (1 )
go func () {
defer wg .Done ()
select {
case <- ctx .Done ():
if err := c .nextConn .SetReadDeadline (veryOld ); err != nil {
errSetDeadline .Store (err )
return
}
<-done
if err := c .nextConn .SetReadDeadline (time .Time {}); err != nil {
errSetDeadline .Store (err )
}
case <- done :
}
}()
n , err := c .nextConn .Read (b )
close (done )
wg .Wait ()
if e := ctx .Err (); e != nil && n == 0 {
err = e
}
if err2 , ok := errSetDeadline .Load ().(error ); ok && err == nil && err2 != nil {
err = err2
}
return n , err
}
func (c *connCtx ) WriteContext (ctx context .Context , b []byte ) (int , error ) {
c .writeMu .Lock ()
defer c .writeMu .Unlock ()
select {
case <- c .closed :
return 0 , ErrClosing
default :
}
done := make (chan struct {})
var wg sync .WaitGroup
var errSetDeadline atomic .Value
wg .Add (1 )
go func () {
defer wg .Done ()
select {
case <- ctx .Done ():
if err := c .nextConn .SetWriteDeadline (veryOld ); err != nil {
errSetDeadline .Store (err )
return
}
<-done
if err := c .nextConn .SetWriteDeadline (time .Time {}); err != nil {
errSetDeadline .Store (err )
}
case <- done :
}
}()
n , err := c .nextConn .Write (b )
close (done )
wg .Wait ()
if e := ctx .Err (); e != nil && n == 0 {
err = e
}
if err2 , ok := errSetDeadline .Load ().(error ); ok && err == nil && err2 != nil {
err = err2
}
return n , err
}
func (c *connCtx ) Close () error {
err := c .nextConn .Close ()
c .closeOnce .Do (func () {
c .writeMu .Lock ()
c .readMu .Lock ()
close (c .closed )
c .readMu .Unlock ()
c .writeMu .Unlock ()
})
return err
}
func (c *connCtx ) LocalAddr () net .Addr {
return c .nextConn .LocalAddr ()
}
func (c *connCtx ) RemoteAddr () net .Addr {
return c .nextConn .RemoteAddr ()
}
func (c *connCtx ) Conn () net .Conn {
return c .nextConn
}
The pages are generated with Golds v0.8.2 . (GOOS=linux GOARCH=amd64)
Golds is a Go 101 project developed by Tapir Liu .
PR and bug reports are welcome and can be submitted to the issue list .
Please follow @zigo_101 (reachable from the left QR code) to get the latest news of Golds .