Calculate per-network traffic (#963)

znc · Aug 21, 2015 · 38fb4cc · 38fb4cc
1 parent 918a8a1
commit 38fb4cc
Show file tree

Hide file tree

Showing 6 changed files with 48 additions and 8 deletions.
diff --git a/include/znc/IRCNetwork.h b/include/znc/IRCNetwork.h
@@ -196,6 +196,12 @@ class CIRCNetwork {
 	unsigned short int GetJoinDelay() const { return m_uJoinDelay; }
 	void SetJoinDelay(unsigned short int uJoinDelay) { m_uJoinDelay = uJoinDelay; }
 
+	unsigned long long BytesRead() const { return m_uBytesRead; }
+	unsigned long long BytesWritten() const { return m_uBytesWritten; }
+
+	void AddBytesRead(unsigned long long u) { m_uBytesRead += u; }
+	void AddBytesWritten(unsigned long long u) { m_uBytesWritten += u; }
+
 	CString ExpandString(const CString& sStr) const;
 	CString& ExpandString(const CString& sStr, CString& sRet) const;
 private:
@@ -245,6 +251,8 @@ class CIRCNetwork {
 	CIRCNetworkJoinTimer* m_pJoinTimer;
 
 	unsigned short int m_uJoinDelay;
+	unsigned long long m_uBytesRead;
+	unsigned long long m_uBytesWritten;
 };
 
 #endif // !ZNC_IRCNETWORK_H
diff --git a/include/znc/User.h b/include/znc/User.h
@@ -187,8 +187,8 @@ class CUser {
 	bool AutoClearQueryBuffer() const;
 	bool IsBeingDeleted() const { return m_bBeingDeleted; }
 	CString GetTimezone() const { return m_sTimezone; }
-	unsigned long long BytesRead() const { return m_uBytesRead; }
-	unsigned long long BytesWritten() const { return m_uBytesWritten; }
+	unsigned long long BytesRead() const;
+	unsigned long long BytesWritten() const;
 	unsigned int JoinTries() const { return m_uMaxJoinTries; }
 	unsigned int MaxJoins() const { return m_uMaxJoins; }
 	CString GetSkinName() const;

diff --git a/src/IRCNetwork.cpp b/src/IRCNetwork.cpp
@@ -143,7 +143,9 @@ CIRCNetwork::CIRCNetwork(CUser *pUser, const CString& sName)
 		  m_NoticeBuffer(),
 		  m_pPingTimer(nullptr),
 		  m_pJoinTimer(nullptr),
-		  m_uJoinDelay(0)
+		  m_uJoinDelay(0),
+		  m_uBytesRead(0),
+		  m_uBytesWritten(0)
 {
 	SetUser(pUser);
 
@@ -301,13 +303,22 @@ CIRCNetwork::~CIRCNetwork() {
 	}
 	m_vQueries.clear();
 
+	CUser* pUser = GetUser();
 	SetUser(nullptr);
 
 	// Make sure we are not in the connection queue
 	CZNC::Get().GetConnectionQueue().remove(this);
 
 	CZNC::Get().GetManager().DelCronByAddr(m_pPingTimer);
 	CZNC::Get().GetManager().DelCronByAddr(m_pJoinTimer);
+
+	if (pUser) {
+		pUser->AddBytesRead(m_uBytesRead);
+		pUser->AddBytesWritten(m_uBytesWritten);
+	} else {
+		CZNC::Get().AddBytesRead(m_uBytesRead);
+		CZNC::Get().AddBytesWritten(m_uBytesWritten);
+	}
 }
 
 void CIRCNetwork::DelServers() {

diff --git a/src/IRCSock.cpp b/src/IRCSock.cpp
@@ -130,8 +130,8 @@ CIRCSock::~CIRCSock() {
 
 	Quit();
 	m_msChans.clear();
-	m_pNetwork->GetUser()->AddBytesRead(GetBytesRead());
-	m_pNetwork->GetUser()->AddBytesWritten(GetBytesWritten());
+	m_pNetwork->AddBytesRead(GetBytesRead());
+	m_pNetwork->AddBytesWritten(GetBytesWritten());
 }
 
 void CIRCSock::Quit(const CString& sQuitMsg) {

diff --git a/src/Socket.cpp b/src/Socket.cpp
@@ -407,14 +407,19 @@ CSocket::CSocket(CModule* pModule, const CString& sHostname, unsigned short uPor
 
 CSocket::~CSocket() {
 	CUser *pUser = nullptr;
+	CIRCNetwork* pNetwork = nullptr;
 
 	// CWebSock could cause us to have a nullptr pointer here
 	if (m_pModule) {
 		pUser = m_pModule->GetUser();
+		pNetwork = m_pModule->GetNetwork();
 		m_pModule->UnlinkSocket(this);
 	}
 
-	if (pUser && m_pModule && (m_pModule->GetType() != CModInfo::GlobalModule)) {
+	if (pNetwork && m_pModule && (m_pModule->GetType() == CModInfo::NetworkModule)) {
+		pNetwork->AddBytesWritten(GetBytesWritten());
+		pNetwork->AddBytesRead(GetBytesRead());
+	} else if (pUser && m_pModule && (m_pModule->GetType() == CModInfo::UserModule)) {
 		pUser->AddBytesWritten(GetBytesWritten());
 		pUser->AddBytesRead(GetBytesRead());
 	} else {

diff --git a/src/User.cpp b/src/User.cpp
@@ -120,8 +120,8 @@ CUser::~CUser() {
 
 	CZNC::Get().GetManager().DelCronByAddr(m_pUserTimer);
 
-	CZNC::Get().AddBytesRead(BytesRead());
-	CZNC::Get().AddBytesWritten(BytesWritten());
+	CZNC::Get().AddBytesRead(m_uBytesRead);
+	CZNC::Get().AddBytesWritten(m_uBytesWritten);
 }
 
 template<class T>
@@ -1277,3 +1277,19 @@ bool CUser::AutoClearQueryBuffer() const { return m_bAutoClearQueryBuffer; }
 CString CUser::GetSkinName() const { return m_sSkinName; }
 const CString& CUser::GetUserPath() const { if (!CFile::Exists(m_sUserPath)) { CDir::MakeDir(m_sUserPath); } return m_sUserPath; }
 // !Getters
+
+unsigned long long CUser::BytesRead() const {
+	unsigned long long uBytes = m_uBytesRead;
+	for (const CIRCNetwork* pNetwork : m_vIRCNetworks) {
+		uBytes += pNetwork->BytesRead();
+	}
+	return uBytes;
+}
+
+unsigned long long CUser::BytesWritten() const {
+	unsigned long long uBytes = m_uBytesWritten;
+	for (const CIRCNetwork* pNetwork : m_vIRCNetworks) {
+		uBytes += pNetwork->BytesWritten();
+	}
+	return uBytes;
+}