[SERVER] Use atomic var for (total)bytesSent counters

Gets rid of the lastBytesSent field as well as the stats lock per
client. Cleaned and split up the messy net_clientsToJson function while
at it.

5 files changed, 62 insertions, 102 deletions

diff --git a/src/server/globals.h b/src/server/globals.h
index c384ac6..84588c3 100644
--- a/src/server/globals.h
+++ b/src/server/globals.h
@@ -127,16 +127,14 @@ struct _dnbd3_image
 struct _dnbd3_client
 {
 #define HOSTNAMELEN (48)
-	uint64_t bytesSent;     // Byte counter for this client. Use statsLock when accessing.
-	uint64_t lastBytesSent; // Byte counter from when we last added to global counter. Use statsLock when accessing.
-	dnbd3_image_t *image;
+	atomic_uint_fast64_t bytesSent;   // Byte counter for this client.
+	dnbd3_image_t *image;             // Image in use by this client, or NULL during handshake
 	int sock;
-	bool isServer;          // true if a server in proxy mode, false if real client
+	bool isServer;                    // true if a server in proxy mode, false if real client
 	dnbd3_host_t host;
-	char hostName[HOSTNAMELEN];
-	pthread_mutex_t sendMutex;
+	char hostName[HOSTNAMELEN];       // inet_ntop version of host
+	pthread_mutex_t sendMutex;        // Held while writing to sock if image is incomplete (since uplink uses socket too)
 	pthread_spinlock_t lock;
-	pthread_spinlock_t statsLock;
 };
 
 // #######################################################
diff --git a/src/server/net.c b/src/server/net.c
index 376af6f..67908cf 100644
--- a/src/server/net.c
+++ b/src/server/net.c
@@ -42,6 +42,7 @@
 #endif
 #include <jansson.h>
 #include <inttypes.h>
+#include <stdatomic.h>
 
 static dnbd3_client_t *_clients[SERVER_MAX_CLIENTS];
 static int _num_clients = 0;
@@ -49,25 +50,13 @@ static pthread_spinlock_t _clients_lock;
 
 static char nullbytes[500];
 
-static uint64_t totalBytesSent = 0;
-static pthread_spinlock_t statisticsSentLock;
+static atomic_uint_fast64_t totalBytesSent = 0;
 
 // Adding and removing clients -- list management
 static bool addToList(dnbd3_client_t *client);
 static void removeFromList(dnbd3_client_t *client);
 static dnbd3_client_t* freeClientStruct(dnbd3_client_t *client);
 
-/**
- * Update global sent stats. Hold client's statsLock when calling.
- */
-void net_updateGlobalSentStatsFromClient(dnbd3_client_t * const client)
-{
-	spin_lock( &statisticsSentLock );
-	totalBytesSent += ( client->bytesSent - client->lastBytesSent );
-	spin_unlock( &statisticsSentLock );
-	client->lastBytesSent = client->bytesSent;
-}
-
 static inline bool recv_request_header(int sock, dnbd3_request_t *request)
 {
 	ssize_t ret, fails = 0;
@@ -154,19 +143,9 @@ static inline bool sendPadding( const int fd, uint32_t bytes )
 	return sock_sendAll( fd, nullbytes, bytes, 2 ) == (ssize_t)bytes;
 }
 
-uint64_t net_getTotalBytesSent()
-{
-	// reads and writes to 64bit ints are not atomic on x86, so let's be safe and use locking
-	spin_lock( &statisticsSentLock );
-	const uint64_t tmp = totalBytesSent;
-	spin_unlock( &statisticsSentLock );
-	return tmp;
-}
-
 void net_init()
 {
 	spin_init( &_clients_lock, PTHREAD_PROCESS_PRIVATE );
-	spin_init( &statisticsSentLock, PTHREAD_PROCESS_PRIVATE );
 }
 
 void* net_handleNewConnection(void *clientPtr)
@@ -208,17 +187,13 @@ void* net_handleNewConnection(void *clientPtr)
 	} while (0);
 	// Fully init client struct
 	spin_init( &client->lock, PTHREAD_PROCESS_PRIVATE );
-	spin_init( &client->statsLock, PTHREAD_PROCESS_PRIVATE );
 	pthread_mutex_init( &client->sendMutex, NULL );
 
 	spin_lock( &client->lock );
 	host_to_string( &client->host, client->hostName, HOSTNAMELEN );
 	client->hostName[HOSTNAMELEN-1] = '\0';
 	spin_unlock( &client->lock );
-	spin_lock( &client->statsLock );
 	client->bytesSent = 0;
-	client->lastBytesSent = 0;
-	spin_unlock( &client->statsLock );
 
 	if ( !addToList( client ) ) {
 		freeClientStruct( client );
@@ -496,10 +471,8 @@ void* net_handleNewConnection(void *clientPtr)
 					}
 				}
 				if ( lock ) pthread_mutex_unlock( &client->sendMutex );
-				spin_lock( &client->statsLock );
 				// Global per-client counter
 				client->bytesSent += request.size; // Increase counter for statistics.
-				spin_unlock( &client->statsLock );
 				break;
 
 			case CMD_GET_SERVERS:
@@ -557,9 +530,8 @@ set_name: ;
 	}
 exit_client_cleanup: ;
 	removeFromList( client );
-	spin_lock( &client->statsLock ); // Make TSAN happy
-	net_updateGlobalSentStatsFromClient( client );
-	spin_unlock( &client->statsLock );
+	// First remove from list, then add to counter to prevent race condition
+	totalBytesSent += client->bytesSent;
 	freeClientStruct( client ); // This will also call image_release on client->image
 	return NULL ;
 fail_preadd: ;
@@ -569,63 +541,67 @@ fail_preadd: ;
 }
 
 /**
- * Get list of all clients and update the global stats counter while we're at it.
- * This method sucks since it has a param that tells it not to generate the list
- * but only update the global counter, which is a horrible relic from refactoring.
- * Hopfully I'll fix it soon by splitting this up or something.
+ * Get list of all clients.
  */
-json_t* net_clientsToJson(const bool fullList)
+struct json_t* net_getListAsJson()
 {
-	json_t *jsonClients = fullList ? json_array() : NULL;
+	json_t *jsonClients = json_array();
 	json_t *clientStats;
-	int i;
 	int imgId;
 	uint64_t bytesSent;
 	char host[HOSTNAMELEN];
 	host[HOSTNAMELEN-1] = '\0';
-	json_int_t counter = 0;
 
 	spin_lock( &_clients_lock );
-	for ( i = 0; i < _num_clients; ++i ) {
-		if ( _clients[i] == NULL ) {
-			continue;
-		}
+	for ( int i = 0; i < _num_clients; ++i ) {
 		dnbd3_client_t * const client = _clients[i];
+		if ( client == NULL || client->image == NULL )
+			continue;
 		spin_lock( &client->lock );
-		spin_unlock( &_clients_lock );
 		// Unlock so we give other threads a chance to access the client list.
 		// We might not get an atomic snapshot of the currently connected clients,
 		// but that doesn't really make a difference anyways.
-		if ( client->image == NULL ) {
-			spin_unlock( &client->lock );
-			imgId = -1;
-		} else {
-			if ( fullList ) {
-				strncpy( host, client->hostName, HOSTNAMELEN - 1 );
-				imgId = client->image->id;
-			} else {
-				counter++;
-			}
-			spin_lock( &client->statsLock );
-			spin_unlock( &client->lock );
-			bytesSent = client->bytesSent;
-			net_updateGlobalSentStatsFromClient( client ); // Do this since we read the totalBytesSent counter later
-			spin_unlock( &client->statsLock );
-		}
-		if ( fullList && imgId != -1 ) {
-			clientStats = json_pack( "{sssisI}",
-					"address", host,
-					"imageId", imgId,
-					"bytesSent", (json_int_t)bytesSent );
-			json_array_append_new( jsonClients, clientStats );
-		}
+		spin_unlock( &_clients_lock );
+		strncpy( host, client->hostName, HOSTNAMELEN - 1 );
+		imgId = client->image->id;
+		bytesSent = client->bytesSent;
+		spin_unlock( &client->lock );
+		clientStats = json_pack( "{sssisI}",
+				"address", host,
+				"imageId", imgId,
+				"bytesSent", (json_int_t)bytesSent );
+		json_array_append_new( jsonClients, clientStats );
 		spin_lock( &_clients_lock );
 	}
 	spin_unlock( &_clients_lock );
-	if ( fullList ) {
-		return jsonClients;
+	return jsonClients;
+}
+
+/**
+ * Get number of clients connected, total bytes sent, or both.
+ * we don't unlock the list while iterating or we might get an
+ * incorrect result if a client is disconnecting while iterating.
+ */
+void net_getStats(int *clientCount, uint64_t *bytesSent)
+{
+	int cc = 0;
+	uint64_t bs = 0;
+
+	spin_lock( &_clients_lock );
+	for ( int i = 0; i < _num_clients; ++i ) {
+		const dnbd3_client_t * const client = _clients[i];
+		if ( client == NULL || client->image == NULL )
+			continue;
+		cc += 1;
+		bs += client->bytesSent;
+	}
+	spin_unlock( &_clients_lock );
+	if ( clientCount != NULL ) {
+		*clientCount = cc;
+	}
+	if ( bytesSent != NULL ) {
+		*bytesSent = totalBytesSent + bs;
 	}
-	return json_integer( counter );
 }
 
 void net_disconnectAll()
@@ -697,8 +673,6 @@ static dnbd3_client_t* freeClientStruct(dnbd3_client_t *client)
 	if ( client->sock != -1 ) close( client->sock );
 	client->sock = -1;
 	pthread_mutex_unlock( &client->sendMutex );
-	spin_lock( &client->statsLock );
-	spin_unlock( &client->statsLock );
 	if ( client->image != NULL ) {
 		spin_lock( &client->image->lock );
 		if ( client->image->uplink != NULL ) uplink_removeClient( client->image->uplink, client );
@@ -707,7 +681,6 @@ static dnbd3_client_t* freeClientStruct(dnbd3_client_t *client)
 	}
 	spin_unlock( &client->lock );
 	spin_destroy( &client->lock );
-	spin_destroy( &client->statsLock );
 	pthread_mutex_destroy( &client->sendMutex );
 	free( client );
 	return NULL ;
diff --git a/src/server/net.h b/src/server/net.h
index 51f4db8..e91395c 100644
--- a/src/server/net.h
+++ b/src/server/net.h
@@ -29,13 +29,9 @@ void net_init();
 
 void* net_handleNewConnection(void *clientPtr);
 
-void net_updateGlobalSentStatsFromClient(dnbd3_client_t * const client);
+struct json_t* net_getListAsJson();
 
-uint64_t net_getTotalBytesSent();
-
-void* net_client_handler(void *client_socket);
-
-struct json_t* net_clientsToJson();
+void net_getStats(int *clientCount, uint64_t *bytesSent);
 
 void net_disconnectAll();
 
diff --git a/src/server/rpc.c b/src/server/rpc.c
index 3d3d665..2e1fca0 100644
--- a/src/server/rpc.c
+++ b/src/server/rpc.c
@@ -289,18 +289,17 @@ static bool handleStatus(int sock, int permissions, struct field *fields, size_t
 	if ( altservers && !(permissions & ACL_ALTSERVERS) ) {
 		return sendReply( sock, "403 Forbidden", "text/plain", "No permission to access altservers", -1, keepAlive );
 	}
-	// Call this first because it will update the total bytes sent counter
-	json_t *jsonClients = NULL;
-	if ( stats || clients ) {
-		jsonClients = net_clientsToJson( clients );
-	}
+
 	json_t *statisticsJson;
 	if ( stats ) {
+		int clientCount;
+		uint64_t bytesSent;
 		const uint64_t bytesReceived = uplink_getTotalBytesReceived();
-		const uint64_t bytesSent = net_getTotalBytesSent();
-		statisticsJson = json_pack( "{sIsIsIsI}",
+		net_getStats( &clientCount, &bytesSent );
+		statisticsJson = json_pack( "{sIsIsisIsI}",
 				"bytesReceived", (json_int_t) bytesReceived,
 				"bytesSent", (json_int_t) bytesSent,
+				"clientCount", clientCount,
 				"uptime", (json_int_t) dnbd3_serverUptime(),
 				"runId", randomRunId );
 	} else {
@@ -313,12 +312,8 @@ static bool handleStatus(int sock, int permissions, struct field *fields, size_t
 		json_object_set_new( statisticsJson, "spaceTotal", json_integer( spaceTotal ) );
 		json_object_set_new( statisticsJson, "spaceFree", json_integer( spaceAvail ) );
 	}
-	if ( jsonClients != NULL ) {
-		if ( clients ) {
-			json_object_set_new( statisticsJson, "clients", jsonClients );
-		} else if ( stats ) {
-			json_object_set_new( statisticsJson, "clientCount", jsonClients );
-		}
+	if ( clients ) {
+		json_object_set_new( statisticsJson, "clients", net_getListAsJson() );
 	}
 	if ( images ) {
 		json_object_set_new( statisticsJson, "images", image_getListAsJson() );
diff --git a/src/server/uplink.c b/src/server/uplink.c
index da63d82..879d08c 100644
--- a/src/server/uplink.c
+++ b/src/server/uplink.c
@@ -802,12 +802,10 @@ static void uplink_handleReceive(dnbd3_connection_t *link)
 						bytesSent = (size_t)sent - sizeof outReply;
 					}
 				}
-				spin_lock( &client->statsLock );
 				pthread_mutex_unlock( &client->sendMutex );
 				if ( bytesSent != 0 ) {
 					client->bytesSent += bytesSent;
 				}
-				spin_unlock( &client->statsLock );
 				spin_lock( &link->queueLock );
 			}
 			if ( req->status == ULR_FREE && i == link->queueLen - 1 ) link->queueLen--;