github.com/micro/go-micro/v2@v2.9.1/util/socket/socket.go (about)

     1  // Package socket provides a pseudo socket
     2  package socket
     3  
     4  import (
     5  	"io"
     6  
     7  	"github.com/micro/go-micro/v2/transport"
     8  )
     9  
    10  // Socket is our pseudo socket for transport.Socket
    11  type Socket struct {
    12  	id string
    13  	// closed
    14  	closed chan bool
    15  	// remote addr
    16  	remote string
    17  	// local addr
    18  	local string
    19  	// send chan
    20  	send chan *transport.Message
    21  	// recv chan
    22  	recv chan *transport.Message
    23  }
    24  
    25  func (s *Socket) SetLocal(l string) {
    26  	s.local = l
    27  }
    28  
    29  func (s *Socket) SetRemote(r string) {
    30  	s.remote = r
    31  }
    32  
    33  // Accept passes a message to the socket which will be processed by the call to Recv
    34  func (s *Socket) Accept(m *transport.Message) error {
    35  	select {
    36  	case s.recv <- m:
    37  		return nil
    38  	case <-s.closed:
    39  		return io.EOF
    40  	}
    41  }
    42  
    43  // Process takes the next message off the send queue created by a call to Send
    44  func (s *Socket) Process(m *transport.Message) error {
    45  	select {
    46  	case msg := <-s.send:
    47  		*m = *msg
    48  	case <-s.closed:
    49  		// see if we need to drain
    50  		select {
    51  		case msg := <-s.send:
    52  			*m = *msg
    53  			return nil
    54  		default:
    55  			return io.EOF
    56  		}
    57  	}
    58  	return nil
    59  }
    60  
    61  func (s *Socket) Remote() string {
    62  	return s.remote
    63  }
    64  
    65  func (s *Socket) Local() string {
    66  	return s.local
    67  }
    68  
    69  func (s *Socket) Send(m *transport.Message) error {
    70  	// send a message
    71  	select {
    72  	case s.send <- m:
    73  	case <-s.closed:
    74  		return io.EOF
    75  	}
    76  
    77  	return nil
    78  }
    79  
    80  func (s *Socket) Recv(m *transport.Message) error {
    81  	// receive a message
    82  	select {
    83  	case msg := <-s.recv:
    84  		// set message
    85  		*m = *msg
    86  	case <-s.closed:
    87  		return io.EOF
    88  	}
    89  
    90  	// return nil
    91  	return nil
    92  }
    93  
    94  // Close closes the socket
    95  func (s *Socket) Close() error {
    96  	select {
    97  	case <-s.closed:
    98  		// no op
    99  	default:
   100  		close(s.closed)
   101  	}
   102  	return nil
   103  }
   104  
   105  // New returns a new pseudo socket which can be used in the place of a transport socket.
   106  // Messages are sent to the socket via Accept and receives from the socket via Process.
   107  // SetLocal/SetRemote should be called before using the socket.
   108  func New(id string) *Socket {
   109  	return &Socket{
   110  		id:     id,
   111  		closed: make(chan bool),
   112  		local:  "local",
   113  		remote: "remote",
   114  		send:   make(chan *transport.Message, 128),
   115  		recv:   make(chan *transport.Message, 128),
   116  	}
   117  }