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 }