gitee.com/liuxuezhan/go-micro-v1.18.0@v1.0.0/util/socket/socket.go (about)

     1  // Package socket provides a pseudo socket
     2  package socket
     3  
     4  import (
     5  	"io"
     6  
     7  	"gitee.com/liuxuezhan/go-micro-v1.18.0/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  	// make copy
    71  	msg := &transport.Message{
    72  		Header: make(map[string]string),
    73  		Body:   make([]byte, len(m.Body)),
    74  	}
    75  
    76  	// copy headers
    77  	for k, v := range m.Header {
    78  		msg.Header[k] = v
    79  	}
    80  
    81  	// copy body
    82  	copy(msg.Body, m.Body)
    83  
    84  	// send a message
    85  	select {
    86  	case s.send <- msg:
    87  	case <-s.closed:
    88  		return io.EOF
    89  	}
    90  
    91  	return nil
    92  }
    93  
    94  func (s *Socket) Recv(m *transport.Message) error {
    95  	// receive a message
    96  	select {
    97  	case msg := <-s.recv:
    98  		// set message
    99  		*m = *msg
   100  	case <-s.closed:
   101  		return io.EOF
   102  	}
   103  
   104  	// return nil
   105  	return nil
   106  }
   107  
   108  // Close closes the socket
   109  func (s *Socket) Close() error {
   110  	select {
   111  	case <-s.closed:
   112  		// no op
   113  	default:
   114  		close(s.closed)
   115  	}
   116  	return nil
   117  }
   118  
   119  // New returns a new pseudo socket which can be used in the place of a transport socket.
   120  // Messages are sent to the socket via Accept and receives from the socket via Process.
   121  // SetLocal/SetRemote should be called before using the socket.
   122  func New(id string) *Socket {
   123  	return &Socket{
   124  		id:     id,
   125  		closed: make(chan bool),
   126  		local:  "local",
   127  		remote: "remote",
   128  		send:   make(chan *transport.Message, 128),
   129  		recv:   make(chan *transport.Message, 128),
   130  	}
   131  }