inmem_socket.go 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. package sockets
  2. import (
  3. "errors"
  4. "net"
  5. "sync"
  6. )
  7. var errClosed = errors.New("use of closed network connection")
  8. // InmemSocket implements net.Listener using in-memory only connections.
  9. type InmemSocket struct {
  10. chConn chan net.Conn
  11. chClose chan struct{}
  12. addr string
  13. mu sync.Mutex
  14. }
  15. // dummyAddr is used to satisfy net.Addr for the in-mem socket
  16. // it is just stored as a string and returns the string for all calls
  17. type dummyAddr string
  18. // NewInmemSocket creates an in-memory only net.Listener
  19. // The addr argument can be any string, but is used to satisfy the `Addr()` part
  20. // of the net.Listener interface
  21. func NewInmemSocket(addr string, bufSize int) *InmemSocket {
  22. return &InmemSocket{
  23. chConn: make(chan net.Conn, bufSize),
  24. chClose: make(chan struct{}),
  25. addr: addr,
  26. }
  27. }
  28. // Addr returns the socket's addr string to satisfy net.Listener
  29. func (s *InmemSocket) Addr() net.Addr {
  30. return dummyAddr(s.addr)
  31. }
  32. // Accept implements the Accept method in the Listener interface; it waits for the next call and returns a generic Conn.
  33. func (s *InmemSocket) Accept() (net.Conn, error) {
  34. select {
  35. case conn := <-s.chConn:
  36. return conn, nil
  37. case <-s.chClose:
  38. return nil, errClosed
  39. }
  40. }
  41. // Close closes the listener. It will be unavailable for use once closed.
  42. func (s *InmemSocket) Close() error {
  43. s.mu.Lock()
  44. defer s.mu.Unlock()
  45. select {
  46. case <-s.chClose:
  47. default:
  48. close(s.chClose)
  49. }
  50. return nil
  51. }
  52. // Dial is used to establish a connection with the in-mem server
  53. func (s *InmemSocket) Dial(network, addr string) (net.Conn, error) {
  54. srvConn, clientConn := net.Pipe()
  55. select {
  56. case s.chConn <- srvConn:
  57. case <-s.chClose:
  58. return nil, errClosed
  59. }
  60. return clientConn, nil
  61. }
  62. // Network returns the addr string, satisfies net.Addr
  63. func (a dummyAddr) Network() string {
  64. return string(a)
  65. }
  66. // String returns the string form
  67. func (a dummyAddr) String() string {
  68. return string(a)
  69. }
  70. // timeoutError is used when there is a timeout with a connection
  71. // this implements the net.Error interface
  72. type timeoutError struct{}
  73. func (e *timeoutError) Error() string { return "i/o timeout" }
  74. func (e *timeoutError) Timeout() bool { return true }
  75. func (e *timeoutError) Temporary() bool { return true }