micro/gateway/listener.go

package gateway

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

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

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

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

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

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

func NewListener(addr net.Addr) *Listener {
	l := &Listener{
		addr: addr,
		ch:   make(chan net.Conn, 5),
	}
	return l
}