github.com/volts-dev/volts@v0.0.0-20240120094013-5e9c65924106/transport/sock.go (about)

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