github.com/keltia/go-ipfs@v0.3.8-0.20150909044612-210793031c63/metrics/stream/metered.go (about)

     1  package meterstream
     2  
     3  import (
     4  	metrics "github.com/ipfs/go-ipfs/metrics"
     5  	inet "github.com/ipfs/go-ipfs/p2p/net"
     6  	peer "github.com/ipfs/go-ipfs/p2p/peer"
     7  	protocol "github.com/ipfs/go-ipfs/p2p/protocol"
     8  )
     9  
    10  type meteredStream struct {
    11  	// keys for accessing metrics data
    12  	protoKey protocol.ID
    13  	peerKey  peer.ID
    14  
    15  	inet.Stream
    16  
    17  	// callbacks for reporting bandwidth usage
    18  	mesSent metrics.StreamMeterCallback
    19  	mesRecv metrics.StreamMeterCallback
    20  }
    21  
    22  func newMeteredStream(base inet.Stream, pid protocol.ID, p peer.ID, recvCB, sentCB metrics.StreamMeterCallback) inet.Stream {
    23  	return &meteredStream{
    24  		Stream:   base,
    25  		mesSent:  sentCB,
    26  		mesRecv:  recvCB,
    27  		protoKey: pid,
    28  		peerKey:  p,
    29  	}
    30  }
    31  
    32  func WrapStream(base inet.Stream, pid protocol.ID, bwc metrics.Reporter) inet.Stream {
    33  	return newMeteredStream(base, pid, base.Conn().RemotePeer(), bwc.LogRecvMessageStream, bwc.LogSentMessageStream)
    34  }
    35  
    36  func (s *meteredStream) Read(b []byte) (int, error) {
    37  	n, err := s.Stream.Read(b)
    38  
    39  	// Log bytes read
    40  	s.mesRecv(int64(n), s.protoKey, s.peerKey)
    41  
    42  	return n, err
    43  }
    44  
    45  func (s *meteredStream) Write(b []byte) (int, error) {
    46  	n, err := s.Stream.Write(b)
    47  
    48  	// Log bytes written
    49  	s.mesSent(int64(n), s.protoKey, s.peerKey)
    50  
    51  	return n, err
    52  }