github.com/GuanceCloud/cliutils@v1.1.21/network/ws/epoll_linux.go (about)

     1  // Unless explicitly stated otherwise all files in this repository are licensed
     2  // under the MIT License.
     3  // This product includes software developed at Guance Cloud (https://www.guance.com/).
     4  // Copyright 2021-present Guance, Inc.
     5  
     6  //go:build linux
     7  // +build linux
     8  
     9  package ws
    10  
    11  import (
    12  	"net"
    13  	"sync"
    14  	"syscall"
    15  
    16  	"golang.org/x/sys/unix"
    17  )
    18  
    19  type epoll struct {
    20  	fd          int
    21  	connections map[int]net.Conn
    22  	lock        *sync.RWMutex
    23  }
    24  
    25  func MkEpoll() (*epoll, error) {
    26  	fd, err := unix.EpollCreate1(0)
    27  	if err != nil {
    28  		return nil, err
    29  	}
    30  	return &epoll{
    31  		fd:          fd,
    32  		lock:        &sync.RWMutex{},
    33  		connections: make(map[int]net.Conn),
    34  	}, nil
    35  }
    36  
    37  func (e *epoll) Add(conn net.Conn) error {
    38  	// Extract file descriptor associated with the connection
    39  	fd := websocketFD(conn)
    40  	err := unix.EpollCtl(e.fd, syscall.EPOLL_CTL_ADD, fd,
    41  		&unix.EpollEvent{
    42  			Events: unix.POLLIN | unix.POLLHUP,
    43  			Fd:     int32(fd),
    44  		})
    45  	if err != nil {
    46  		return err
    47  	}
    48  	e.lock.Lock()
    49  	defer e.lock.Unlock()
    50  	e.connections[fd] = conn
    51  	return nil
    52  }
    53  
    54  func (e *epoll) Remove(conn net.Conn) error {
    55  	fd := websocketFD(conn)
    56  	err := unix.EpollCtl(e.fd, syscall.EPOLL_CTL_DEL, fd, nil)
    57  	if err != nil {
    58  		return err
    59  	}
    60  	e.lock.Lock()
    61  	defer e.lock.Unlock()
    62  	delete(e.connections, fd)
    63  	if len(e.connections)%100 == 0 {
    64  		l.Debugf("Total number of connections: %v", len(e.connections))
    65  	}
    66  	return nil
    67  }
    68  
    69  func (e *epoll) Wait(cnt int) ([]net.Conn, error) {
    70  	events := make([]unix.EpollEvent, cnt)
    71  	n, err := unix.EpollWait(e.fd, events, cnt)
    72  	if err != nil {
    73  		return nil, err
    74  	}
    75  	e.lock.RLock()
    76  	defer e.lock.RUnlock()
    77  	var connections []net.Conn
    78  	for i := 0; i < n; i++ {
    79  		conn := e.connections[int(events[i].Fd)]
    80  		connections = append(connections, conn)
    81  	}
    82  	return connections, nil
    83  }
    84  
    85  func (e *epoll) Close() error {
    86  	for _, c := range e.connections {
    87  		l.Debugf("remove cli %s", c.RemoteAddr().String())
    88  
    89  		if err := c.Close(); err != nil {
    90  			l.Errorf("c.Close(): %s", err.Error())
    91  			return err
    92  		}
    93  
    94  		// if err := e.Remove(c); err != nil {
    95  		//	l.Errorf("e.Remove(): %s", err.Error())
    96  		//	return err
    97  		//}
    98  	}
    99  
   100  	if err := unix.Close(e.fd); err != nil {
   101  		l.Errorf("unix.Close(): %s", err.Error())
   102  	}
   103  
   104  	return nil
   105  }