github.com/codingeasygo/util@v0.0.0-20231206062002-1ce2f004b7d9/xio/listener.go (about)

     1  package xio
     2  
     3  import (
     4  	"net"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  type TimeoutListener struct {
    10  	net.Listener
    11  	Delay   time.Duration
    12  	Timeout time.Duration
    13  	waiter  sync.WaitGroup
    14  	running bool
    15  	allConn map[*timeoutConn]int
    16  	allLock sync.RWMutex
    17  }
    18  
    19  func NewTimeoutListener(ln net.Listener, timeout time.Duration) (listener *TimeoutListener) {
    20  	listener = &TimeoutListener{
    21  		Listener: ln,
    22  		Delay:    time.Second,
    23  		Timeout:  timeout,
    24  		waiter:   sync.WaitGroup{},
    25  		running:  true,
    26  		allConn:  map[*timeoutConn]int{},
    27  		allLock:  sync.RWMutex{},
    28  	}
    29  	listener.waiter.Add(1)
    30  	go listener.runTimeout()
    31  	return
    32  }
    33  
    34  func (t *TimeoutListener) runTimeout() {
    35  	for t.running {
    36  		now := time.Now()
    37  		t.allLock.Lock()
    38  		for c := range t.allConn {
    39  			if now.Sub(c.latest) >= t.Timeout {
    40  				c.rawClose()
    41  				delete(t.allConn, c)
    42  			}
    43  		}
    44  		t.allLock.Unlock()
    45  		time.Sleep(t.Delay)
    46  	}
    47  	t.waiter.Done()
    48  }
    49  
    50  func (t *TimeoutListener) Accept() (conn net.Conn, err error) {
    51  	conn, err = t.Listener.Accept()
    52  	if err == nil {
    53  		t.allLock.Lock()
    54  		c := &timeoutConn{Conn: conn, listener: t, latest: time.Now()}
    55  		t.allConn[c] = 1
    56  		t.allLock.Unlock()
    57  		conn = c
    58  	}
    59  	return
    60  }
    61  
    62  func (t *TimeoutListener) Close() (err error) {
    63  	err = t.Listener.Close()
    64  	t.running = false
    65  	t.waiter.Wait()
    66  	return
    67  }
    68  
    69  func (t *TimeoutListener) closeConn(c *timeoutConn) {
    70  	t.allLock.Lock()
    71  	delete(t.allConn, c)
    72  	t.allLock.Unlock()
    73  }
    74  
    75  type timeoutConn struct {
    76  	net.Conn
    77  	listener *TimeoutListener
    78  	latest   time.Time
    79  }
    80  
    81  func (t *timeoutConn) Read(p []byte) (n int, err error) {
    82  	t.latest = time.Now()
    83  	n, err = t.Conn.Read(p)
    84  	return
    85  }
    86  
    87  func (t *timeoutConn) Write(p []byte) (n int, err error) {
    88  	t.latest = time.Now()
    89  	n, err = t.Conn.Write(p)
    90  	return
    91  }
    92  
    93  func (t *timeoutConn) rawClose() (err error) {
    94  	t.latest = time.Now()
    95  	err = t.Conn.Close()
    96  	return
    97  }
    98  
    99  func (t *timeoutConn) Close() (err error) {
   100  	err = t.rawClose()
   101  	t.listener.closeConn(t)
   102  	return
   103  }