Defer counters

wrappers/mtproto_abridged.go

@@ -28,6 +28,10 @@ type MTProtoAbridged struct {
 }
 
 func (m *MTProtoAbridged) Read() ([]byte, error) {
+	defer func() {
+		m.readCounter++
+	}()
+
 	m.logger.Debugw("Read packet",
 		"simple_ack", m.opts.ReadHacks.SimpleAck,
 		"quick_ack", m.opts.ReadHacks.QuickAck,
@@ -75,16 +79,18 @@ func (m *MTProtoAbridged) Read() ([]byte, error) {
 
 	buf.Reset()
 	buf.Grow(int(msgLength32))
-
 	if _, err := io.CopyN(buf, m.conn, int64(msgLength32)); err != nil {
 		return nil, errors.Annotate(err, "Cannot read message")
 	}
-	m.readCounter++
 
 	return buf.Bytes(), nil
 }
 
 func (m *MTProtoAbridged) Write(p []byte) (int, error) {
+	defer func() {
+		m.writeCounter++
+	}()
+
 	m.logger.Debugw("Write packet",
 		"length", len(p),
 		"simple_ack", m.opts.WriteHacks.SimpleAck,
@@ -97,7 +103,6 @@ func (m *MTProtoAbridged) Write(p []byte) (int, error) {
 	}
 
 	if m.opts.WriteHacks.SimpleAck {
-		m.writeCounter++
 		return m.conn.Write(utils.ReverseBytes(p))
 	}
 
@@ -105,8 +110,6 @@ func (m *MTProtoAbridged) Write(p []byte) (int, error) {
 	switch {
 	case packetLength < mtprotoAbridgedSmallPacketLength:
 		newData := append([]byte{byte(packetLength)}, p...)
-
-		m.writeCounter++
 		return m.conn.Write(newData)
 
 	case packetLength < mtprotoAbridgedLargePacketLength:
@@ -119,7 +122,6 @@ func (m *MTProtoAbridged) Write(p []byte) (int, error) {
 		buf.Write(length24[:])
 		buf.Write(p)
 
-		m.writeCounter++
 		return m.conn.Write(buf.Bytes())
 	}
 

wrappers/mtproto_intermediate.go

@@ -24,6 +24,10 @@ type MTProtoIntermediate struct {
 }
 
 func (m *MTProtoIntermediate) Read() ([]byte, error) {
+	defer func() {
+		m.readCounter++
+	}()
+
 	m.logger.Debugw("Read packet",
 		"simple_ack", m.opts.ReadHacks.SimpleAck,
 		"quick_ack", m.opts.ReadHacks.QuickAck,
@@ -59,18 +63,20 @@ func (m *MTProtoIntermediate) Read() ([]byte, error) {
 	if length%4 != 0 {
 		length -= length % 4
 	}
-	m.readCounter++
 
 	return buf.Bytes()[:length], nil
 }
 
 func (m *MTProtoIntermediate) Write(p []byte) (int, error) {
+	defer func() {
+		m.writeCounter++
+	}()
+
 	m.logger.Debugw("Write packet",
 		"simple_ack", m.opts.WriteHacks.SimpleAck,
 		"quick_ack", m.opts.WriteHacks.QuickAck,
 		"counter", m.writeCounter,
 	)
-	m.writeCounter++
 
 	if m.opts.ReadHacks.SimpleAck {
 		return m.conn.Write(p)

wrappers/mtproto_proxy.go

@@ -22,6 +22,10 @@ type MTProtoProxy struct {
 }
 
 func (m *MTProtoProxy) Read() ([]byte, error) {
+	defer func() {
+		m.readCounter++
+	}()
+
 	m.logger.Debugw("Read packet",
 		"counter", m.readCounter,
 		"simple_ack", m.req.Options.WriteHacks.SimpleAck,
@@ -32,9 +36,6 @@ func (m *MTProtoProxy) Read() ([]byte, error) {
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot read packet")
 	}
-	defer func() {
-		m.readCounter++
-	}()
 
 	m.logger.Debugw("Read packet length",
 		"counter", m.readCounter,
@@ -95,13 +96,16 @@ func (m *MTProtoProxy) readCloseExt(data []byte) ([]byte, error) {
 }
 
 func (m *MTProtoProxy) Write(p []byte) (int, error) {
+	defer func() {
+		m.writeCounter++
+	}()
+
 	m.logger.Debugw("Write packet",
 		"length", len(p),
 		"counter", m.writeCounter,
 		"simple_ack", m.req.Options.ReadHacks.SimpleAck,
 		"quick_ack", m.req.Options.ReadHacks.QuickAck,
 	)
-	m.writeCounter++
 
 	header := m.req.MakeHeader(p)
 	if ce := m.logger.Desugar().Check(zap.DebugLevel, "RPC_PROXY_REQ header"); ce != nil {