Go1.9安全的限制net.Listener接入连接数

package main

import (
    "fmt"
    "net"
    "sync"
    "sync/atomic"
)

func main() {
    lis, err := net.Listen("tcp", ":1789")
    if err != nil {
        return
    }

    lis = LimitListener(lis, 1)
}

func LimitListener(l net.Listener, n int) net.Listener {
    return &limitListener{l, 0, make(chan struct{}, n)}
}

type limitListener struct {
    net.Listener
    closed int32
    limit  chan struct{}
}

func (l *limitListener) acquire() { l.limit <- struct{}{} }

func (l *limitListener) release() { <-l.limit }

func (l *limitListener) shutdown() bool { return atomic.LoadInt32(&l.closed) != 0 }

func (l *limitListener) Accept() (net.Conn, error) {
    if l.shutdown() {
        return nil, fmt.Errorf("accept tcp %s: use of closed network connection", l.Listener.Addr().String())
    }
    l.acquire()
    c, err := l.Listener.Accept()
    if err != nil {
        l.release()
        return nil, err
    }
    return &limitListenerConn{Conn: c, release: l.release}, nil
}

func (l *limitListener) Close() error {
    if l.shutdown() {
        return nil
    }
    atomic.StoreInt32(&l.closed, 1)
    close(l.limit)
    return l.Listener.Close()
}

type limitListenerConn struct {
    net.Conn
    releaseOnce sync.Once
    release     func()
}

func (l *limitListenerConn) Close() error {
    err := l.Conn.Close()
    l.releaseOnce.Do(l.release)
    return err
}

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值