kos/entry/listener.go

package entry

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

type Listener struct {
	addr     net.Addr
	ch       chan net.Conn
	exitFlag int32
	exitChan chan struct{}
}

func (l *Listener) Receive(conn net.Conn) {
	select {
	case l.ch <- conn:
	case <-l.exitChan:
	}
}

func (l *Listener) Accept() (net.Conn, error) {
	select {
	case conn, ok := <-l.ch:
		if ok {
			return conn, nil
		} else {
			return nil, io.ErrClosedPipe
		}
	case <-l.exitChan:
		return nil, io.ErrClosedPipe
	}
}

func (l *Listener) Close() error {
	if atomic.CompareAndSwapInt32(&l.exitFlag, 0, 1) {
		close(l.exitChan)
	}
	return nil
}

func (l *Listener) Addr() net.Addr {
	return l.addr
}

func newListener(addr net.Addr) *Listener {
	return &Listener{
		addr:     addr,
		ch:       make(chan net.Conn, 10),
		exitChan: make(chan struct{}),
	}
}