github.com/searKing/golang/go@v1.2.117/net/notify_listener.go (about)

     1  // Copyright 2020 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package net
     6  
     7  import (
     8  	"net"
     9  	"sync"
    10  )
    11  
    12  var _ net.Listener = (*NotifyListener)(nil)
    13  
    14  type notifyAddr string
    15  
    16  func (notifyAddr) Network() string  { return "notify+net" }
    17  func (f notifyAddr) String() string { return string(f) }
    18  
    19  type NotifyListener struct {
    20  	C chan net.Conn
    21  
    22  	mu       sync.Mutex
    23  	doneChan chan struct{}
    24  }
    25  
    26  // NewNotifyListener creates a new Listener that will recv
    27  // conns on its channel.
    28  func NewNotifyListener() *NotifyListener {
    29  	c := make(chan net.Conn)
    30  
    31  	l := &NotifyListener{
    32  		C: c,
    33  	}
    34  	return l
    35  }
    36  
    37  // Addr returns the listener's network address.
    38  func (l *NotifyListener) Addr() net.Addr {
    39  	return notifyAddr("notify_listener")
    40  }
    41  
    42  func (l *NotifyListener) Accept() (net.Conn, error) {
    43  	select {
    44  	case <-l.getDoneChan():
    45  		return nil, ErrListenerClosed
    46  	case c, ok := <-l.C:
    47  		if !ok {
    48  			// Already closed. Don't Accept again.
    49  			return nil, ErrListenerClosed
    50  		}
    51  		return c, nil
    52  	}
    53  }
    54  
    55  // DoneC returns whether this listener has been closed
    56  // for multi producers of C
    57  func (l *NotifyListener) DoneC() <-chan struct{} {
    58  	return l.getDoneChan()
    59  }
    60  
    61  // Close prevents the NotifyListener from firing.
    62  // To ensure the channel is empty after a call to Close, check the
    63  // return value and drain the channel.
    64  func (l *NotifyListener) Close() error {
    65  	l.mu.Lock()
    66  	defer l.mu.Unlock()
    67  	l.closeDoneChanLocked()
    68  
    69  	return l.drainChanLocked()
    70  }
    71  
    72  func (l *NotifyListener) drainChanLocked() error {
    73  	var err error
    74  	// Drain the connections enqueued for the listener.
    75  L:
    76  	for {
    77  		select {
    78  		case c, ok := <-l.C:
    79  			if !ok {
    80  				// Already closed. Don't close again.
    81  				return nil
    82  			}
    83  			if c == nil {
    84  				continue
    85  			}
    86  			if cerr := c.Close(); cerr != nil && err == nil {
    87  				err = cerr
    88  			}
    89  		default:
    90  			break L
    91  		}
    92  	}
    93  	return err
    94  }
    95  
    96  func (l *NotifyListener) getDoneChan() <-chan struct{} {
    97  	l.mu.Lock()
    98  	defer l.mu.Unlock()
    99  	return l.getDoneChanLocked()
   100  }
   101  
   102  func (l *NotifyListener) getDoneChanLocked() chan struct{} {
   103  	if l.doneChan == nil {
   104  		l.doneChan = make(chan struct{})
   105  	}
   106  	return l.doneChan
   107  }
   108  
   109  func (l *NotifyListener) closeDoneChanLocked() {
   110  	ch := l.getDoneChanLocked()
   111  	select {
   112  	case <-ch:
   113  		// Already closed. Don't close again.
   114  	default:
   115  		// Safe to close here. We're the only closer, guarded
   116  		// by s.mu.
   117  		close(ch)
   118  	}
   119  }