Signed-off-by: Rafael David Tinoco <tinhocas@gmail.com>
diff --git a/include/linux/task_io_accounting.h b/include/linux/task_io_accounting.h
index bdf855c..bd45b92 100644
--- a/include/linux/task_io_accounting.h
+++ b/include/linux/task_io_accounting.h
@@ -41,5 +41,12 @@ struct task_io_accounting {
 	 * information loss in doing that.
 	 */
 	u64 cancelled_write_bytes;
+
+	/* The number of bytes which this task has read from a socket */
+	u64 read_net_bytes;
+
+	/* The number of bytes which this task has written to a socket */
+	u64 write_net_bytes;
+
 #endif /* CONFIG_TASK_IO_ACCOUNTING */
 };
diff --git a/include/linux/task_io_accounting_ops.h b/include/linux/task_io_accounting_ops.h
index 4d090f9..ee8416f 100644
--- a/include/linux/task_io_accounting_ops.h
+++ b/include/linux/task_io_accounting_ops.h
@@ -12,6 +12,11 @@ static inline void task_io_account_read(size_t bytes)
 	current->ioac.read_bytes += bytes;
 }
 
+static inline void task_io_account_read_net(size_t bytes)
+{
+	current->ioac.read_net_bytes += bytes;
+}
+
 /*
  * We approximate number of blocks, because we account bytes only.
  * A 'block' is 512 bytes
@@ -26,6 +31,11 @@ static inline void task_io_account_write(size_t bytes)
 	current->ioac.write_bytes += bytes;
 }
 
+static inline void task_io_account_write_net(size_t bytes)
+{
+	current->ioac.write_net_bytes += bytes;
+}
+
 /*
  * We approximate number of blocks, because we account bytes only.
  * A 'block' is 512 bytes
@@ -59,6 +69,10 @@ static inline void task_io_account_read(size_t bytes)
 {
 }
 
+static inline void task_io_account_read_net(size_t bytes)
+{
+}
+
 static inline unsigned long task_io_get_inblock(const struct task_struct *p)
 {
 	return 0;
@@ -68,6 +82,10 @@ static inline void task_io_account_write(size_t bytes)
 {
 }
 
+static inline void task_io_account_write_net(size_t bytes)
+{
+}
+
 static inline unsigned long task_io_get_oublock(const struct task_struct *p)
 {
 	return 0;
diff --git a/include/linux/taskstats.h b/include/linux/taskstats.h
index 341dddb..b0c5990 100644
--- a/include/linux/taskstats.h
+++ b/include/linux/taskstats.h
@@ -163,6 +163,10 @@ struct taskstats {
 	/* Delay waiting for memory reclaim */
 	__u64	freepages_count;
 	__u64	freepages_delay_total;
+
+	/* Per-task network I/O accounting */
+	__u64   read_net_bytes;         /* bytes of socket read I/O */
+	__u64   write_net_bytes;        /* bytes of socket write I/O */
 };
 
 
diff --git a/kernel/tsacct.c b/kernel/tsacct.c
index 00d59d0..b279e69 100644
--- a/kernel/tsacct.c
+++ b/kernel/tsacct.c
@@ -104,10 +104,14 @@ void xacct_add_tsk(struct taskstats *stats, struct task_struct *p)
 	stats->read_bytes	= p->ioac.read_bytes;
 	stats->write_bytes	= p->ioac.write_bytes;
 	stats->cancelled_write_bytes = p->ioac.cancelled_write_bytes;
+	stats->read_net_bytes	= p->ioac.read_net_bytes;
+	stats->write_net_bytes	= p->ioac.write_net_bytes;
 #else
 	stats->read_bytes	= 0;
 	stats->write_bytes	= 0;
 	stats->cancelled_write_bytes = 0;
+	stats->read_net_bytes	= 0;
+	stats->write_net_bytes	= 0;
 #endif
 }
 #undef KB
diff --git a/net/socket.c b/net/socket.c
index 769c386..dd7dbb6 100644
--- a/net/socket.c
+++ b/net/socket.c
@@ -87,6 +87,7 @@
 #include <linux/wireless.h>
 #include <linux/nsproxy.h>
 #include <linux/magic.h>
+#include <linux/task_io_accounting_ops.h>
 
 #include <asm/uaccess.h>
 #include <asm/unistd.h>
@@ -538,6 +539,7 @@ EXPORT_SYMBOL(sock_tx_timestamp);
 static inline int __sock_sendmsg(struct kiocb *iocb, struct socket *sock,
 				 struct msghdr *msg, size_t size)
 {
+	int ret;
 	struct sock_iocb *si = kiocb_to_siocb(iocb);
 	int err;
 
@@ -550,7 +552,12 @@ static inline int __sock_sendmsg(struct kiocb *iocb, struct socket *sock,
 	if (err)
 		return err;
 
-	return sock->ops->sendmsg(iocb, sock, msg, size);
+	ret = sock->ops->sendmsg(iocb, sock, msg, size);
+
+	if (ret > 0)
+		task_io_account_write_net(ret);
+
+	return ret;
 }
 
 int sock_sendmsg(struct socket *sock, struct msghdr *msg, size_t size)
@@ -666,6 +673,7 @@ EXPORT_SYMBOL_GPL(sock_recv_ts_and_drops);
 static inline int __sock_recvmsg_nosec(struct kiocb *iocb, struct socket *sock,
 				       struct msghdr *msg, size_t size, int flags)
 {
+	int ret = 0;
 	struct sock_iocb *si = kiocb_to_siocb(iocb);
 
 	si->sock = sock;
@@ -674,7 +682,12 @@ static inline int __sock_recvmsg_nosec(struct kiocb *iocb, struct socket *sock,
 	si->size = size;
 	si->flags = flags;
 
-	return sock->ops->recvmsg(iocb, sock, msg, size, flags);
+	ret = sock->ops->recvmsg(iocb, sock, msg, size, flags);
+
+	if (ret > 0)
+		task_io_account_read_net(ret);
+
+	return ret;
 }
 
 static inline int __sock_recvmsg(struct kiocb *iocb, struct socket *sock,
