github.com/xmidt-org/webpa-common@v1.11.9/device/connection.go (about)

     1  package device
     2  
     3  import (
     4  	"io"
     5  	"time"
     6  
     7  	"github.com/gorilla/websocket"
     8  	"github.com/xmidt-org/webpa-common/xmetrics"
     9  )
    10  
    11  // Reader represents the read behavior of a device connection
    12  type Reader interface {
    13  	ReadMessage() (int, []byte, error)
    14  	SetReadDeadline(time.Time) error
    15  	SetPongHandler(func(string) error)
    16  }
    17  
    18  // ReadCloser adds io.Closer behavior to Reader
    19  type ReadCloser interface {
    20  	io.Closer
    21  	Reader
    22  }
    23  
    24  // Writer represents the write behavior of a device connection
    25  type Writer interface {
    26  	WriteMessage(int, []byte) error
    27  	WritePreparedMessage(*websocket.PreparedMessage) error
    28  	SetWriteDeadline(time.Time) error
    29  }
    30  
    31  // WriteCloser adds io.Closer behavior to Writer
    32  type WriteCloser interface {
    33  	io.Closer
    34  	Writer
    35  }
    36  
    37  // Connection describes the set of behaviors for device connections used by this package.
    38  type Connection interface {
    39  	io.Closer
    40  	Reader
    41  	Writer
    42  }
    43  
    44  func zeroDeadline() time.Time {
    45  	return time.Time{}
    46  }
    47  
    48  // NewDeadline creates a deadline closure given a timeout and a now function.
    49  // If timeout is nonpositive, the return closure always returns zero time.
    50  // If now is nil (and timeout is positive), then time.Now is used.
    51  func NewDeadline(timeout time.Duration, now func() time.Time) func() time.Time {
    52  	if timeout > 0 {
    53  		if now == nil {
    54  			now = time.Now
    55  		}
    56  
    57  		return func() time.Time {
    58  			return now().Add(timeout)
    59  		}
    60  	}
    61  
    62  	return zeroDeadline
    63  }
    64  
    65  // NewPinger creates a ping closure for the given connection.  Internally, a prepared message is created using the
    66  // supplied data, and the given counter is incremented for each successful update of the write deadline.
    67  func NewPinger(w Writer, pings xmetrics.Incrementer, data []byte, deadline func() time.Time) (func() error, error) {
    68  	pm, err := websocket.NewPreparedMessage(websocket.PingMessage, data)
    69  	if err != nil {
    70  		return nil, err
    71  	}
    72  
    73  	return func() error {
    74  		err := w.SetWriteDeadline(deadline())
    75  		if err != nil {
    76  			return err
    77  		}
    78  
    79  		err = w.WritePreparedMessage(pm)
    80  		if err != nil {
    81  			return err
    82  		}
    83  
    84  		// only incrememt when the complete ping operation was successful
    85  		pings.Inc()
    86  		return nil
    87  	}, nil
    88  }
    89  
    90  // SetPongHandler establishes an instrumented pong handler for the given connection that enforces
    91  // the given read timeout.
    92  func SetPongHandler(r Reader, pongs xmetrics.Incrementer, deadline func() time.Time) {
    93  	r.SetPongHandler(func(_ string) error {
    94  		// increment up front, as this function is only called when a pong is actually received
    95  		pongs.Inc()
    96  		return r.SetReadDeadline(deadline())
    97  	})
    98  }
    99  
   100  type instrumentedReader struct {
   101  	ReadCloser
   102  	statistics Statistics
   103  }
   104  
   105  func (ir *instrumentedReader) ReadMessage() (int, []byte, error) {
   106  	messageType, data, err := ir.ReadCloser.ReadMessage()
   107  	if err == nil {
   108  		ir.statistics.AddBytesReceived(len(data))
   109  		ir.statistics.AddMessagesReceived(1)
   110  	}
   111  
   112  	return messageType, data, err
   113  }
   114  
   115  func InstrumentReader(r ReadCloser, s Statistics) ReadCloser {
   116  	return &instrumentedReader{r, s}
   117  }
   118  
   119  type instrumentedWriter struct {
   120  	WriteCloser
   121  	statistics Statistics
   122  }
   123  
   124  func (iw *instrumentedWriter) WriteMessage(messageType int, data []byte) error {
   125  	err := iw.WriteCloser.WriteMessage(messageType, data)
   126  	if err != nil {
   127  		return err
   128  	}
   129  
   130  	iw.statistics.AddBytesSent(len(data))
   131  	iw.statistics.AddMessagesSent(1)
   132  	return nil
   133  }
   134  
   135  func (iw *instrumentedWriter) WritePreparedMessage(pm *websocket.PreparedMessage) error {
   136  	err := iw.WriteCloser.WritePreparedMessage(pm)
   137  	if err != nil {
   138  		return err
   139  	}
   140  
   141  	// TODO: There isn't any way to obtain the length of a prepared message, so there's not a way to instrument it
   142  	// at the moment
   143  	iw.statistics.AddMessagesSent(1)
   144  	return nil
   145  }
   146  
   147  func InstrumentWriter(w WriteCloser, s Statistics) WriteCloser {
   148  	return &instrumentedWriter{w, s}
   149  }