Improve locking during heavy TCP forking

src/daemon.c

@@ -218,7 +218,9 @@ void cleanup(const int ret)
 		terminate_threads();
 
 		// Close database connection
+		lock_shm();
 		gravityDB_close();
+		unlock_shm();
 
 		// Close sockets and delete Unix socket file handle
 		close_telnet_socket();

src/database/aliasclients.c

@@ -131,6 +131,13 @@ bool import_aliasclients(sqlite3 *db)
 		const int clientID = findClientID(aliasclient_str, false, true);
 
 		clientsData *client = getClient(clientID, true);
+		if(client == NULL)
+		{
+			free(aliasclient_str);
+			continue;
+		}
+
+		// Set client flags
 		client->flags.new = false;
 
 		// Reset counter

src/database/common.c

@@ -344,7 +344,9 @@ void db_init(void)
 		dbversion = db_get_int(db, DB_VERSION);
 	}
 
+	lock_shm();
 	import_aliasclients(db);
+	unlock_shm();
 
 	// Close database to prevent having it opened all time
 	// We already closed the database when we returned earlier

src/database/gravity-db.c

@@ -934,7 +934,8 @@ void gravityDB_close(void)
 	for(int clientID = 0; clientID < counters->clients; clientID++)
 	{
 		clientsData *client = getClient(clientID, true);
-		gravityDB_finalize_client_statements(client);
+		if(client != NULL)
+			gravityDB_finalize_client_statements(client);
 	}
 
 	// Free allocated memory for vectors of prepared client statements

src/database/network-table.c

@@ -1227,8 +1227,6 @@ void parse_neighbor_cache(sqlite3* db)
 
 		// Get hostname of this client if the client is known
 		char *hostname = NULL;
-		// Get client pointer
-		clientsData *client = NULL;
 		bool client_valid = false;
 		time_t lastQuery = 0;
 		unsigned int numQueries = 0;
@@ -1238,7 +1236,7 @@ void parse_neighbor_cache(sqlite3* db)
 		if(clientID >= 0)
 		{
 			client_status[clientID] = CLIENT_ARP_COMPLETE;
-			client = getClient(clientID, true);
+			clientsData *client = getClient(clientID, true);
 			if(client != NULL)
 			{
 				client_valid = true;
@@ -1275,15 +1273,16 @@ void parse_neighbor_cache(sqlite3* db)
 				// Create new record (INSERT)
 				insert_netDB_device(db, hwaddr, now, lastQuery, numQueries, macVendor);
 
+				lock_shm();
+				clientsData *client = getClient(clientID, true);
 				if(client != NULL)
 				{
 					// Reacquire client pointer (if may have changed when unlocking above)
-					lock_shm();
 					client = getClient(clientID, true);
 					// Reset client ARP counter (we stored the entry in the database)
 					client->numQueriesARP = 0;
-					unlock_shm();
 				}
+				unlock_shm();
 
 				// Obtain ID which was given to this new entry
 				dbID = sqlite3_last_insert_rowid(db);
@@ -1330,23 +1329,21 @@ void parse_neighbor_cache(sqlite3* db)
 			if(rc != SQLITE_OK)
 				break;
 
-			// Reacquire client pointer (if may have changed when unlocking above)
-			client = getClient(clientID, true);
 
 			// Update number of queries if applicable
 			rc = update_netDB_numQueries(db, dbID, numQueries);
 			if(rc != SQLITE_OK)
 				break;
 
+			lock_shm();
+			// Acquire client pointer
+			clientsData *client = getClient(clientID, true);
 			if(client != NULL)
 			{
-				// Reacquire client pointer (if may have changed when unlocking above)
-				lock_shm();
-				client = getClient(clientID, true);
 				// Reset client ARP counter (we stored the entry in the database)
 				client->numQueriesARP = 0;
-				unlock_shm();
 			}
+			unlock_shm();
 
 			// Update hostname if available
 			rc = update_netDB_name(db, ip, hostname);

src/database/query-table.c

@@ -369,6 +369,9 @@ void DB_read_queries(void)
 		return;
 	}
 
+	// Lock shared memory
+	lock_shm();
+
 	// Loop through returned database rows
 	while((rc = sqlite3_step(stmt)) == SQLITE_ROW)
 	{
@@ -430,8 +433,8 @@ void DB_read_queries(void)
 			continue;
 		}
 
-		// Lock shared memory
-		lock_shm();
+		// Ensure we have enough shared memory available for new data
+		shm_ensure_size();
 
 		const char *buffer = NULL;
 		int upstreamID = -1; // Default if not forwarded
@@ -591,9 +594,9 @@ void DB_read_queries(void)
 				logg("Warning: Found unknown status %i in long term database!", status);
 				break;
 		}
-
-		unlock_shm();
 	}
+
+	unlock_shm();
 	logg("Imported %i queries from the long-term database", counters->queries);
 
 	// Update lastdbindex so that the next call to DB_save_queries()

src/dnsmasq_interface.c

@@ -237,14 +237,20 @@ static bool _FTL_check_blocking(int queryID, int domainID, int clientID, const c
 	}
 
 	// Get query, domain and client pointers
-	queriesData* query  = getQuery(queryID,   true);
-	domainsData* domain = getDomain(domainID, true);
-	clientsData* client = getClient(clientID, true);
+	queriesData *query  = getQuery(queryID, true);
+	domainsData *domain = getDomain(domainID, true);
+	clientsData *client = getClient(clientID, true);
+	if(query == NULL || domain == NULL || client == NULL)
+	{
+		logg("Error: No memory available, skipping query analysis");
+		return false;
+	}
+
+	// Get cache pointer
 	unsigned int cacheID = findCacheID(domainID, clientID, query->type);
 	DNSCacheData *dns_cache = getDNSCache(cacheID, true);
-	if(query == NULL || domain == NULL || client == NULL || dns_cache == NULL)
+	if(dns_cache == NULL)
 	{
-		// Encountered memory error, skip query
 		logg("WARN: No memory available, skipping query analysis");
 		return false;
 	}
@@ -362,7 +368,7 @@ static bool _FTL_check_blocking(int queryID, int domainID, int clientID, const c
 		return false;
 	}
 
-	// Make a local copy of the domain string. The  string memory may get
+	// Make a local copy of the domain string. The string memory may get
 	// reorganized in the following. We cannot expect domainstr to remain
 	// valid for all time.
 	domainstr = strdup(domainstr);
@@ -487,6 +493,12 @@ bool _FTL_CNAME(const char *domain, const struct crec *cpp, const int id, const
 	{
 		// Increase blocked count of parent domain
 		domainsData* parent_domain = getDomain(parent_domainID, true);
+		if(parent_domain == NULL)
+		{
+			// Memory error, return
+			unlock_shm();
+			return false;
+		}
 		parent_domain->blockedcount++;
 
 		// Store query response as CNAME type
@@ -2044,8 +2056,14 @@ void FTL_TCP_worker_terminating(bool finished)
 		return;
 	}
 
+	// First check if we already locked before. This can happen when a fork
+	// is running into a timeout while it is still processing something and
+	// still holding a lock.
+	if(!is_our_lock())
+		lock_shm();
 	// Close dedicated database connections of this fork
 	gravityDB_close();
+	unlock_shm();
 }
 
 // Called when a (forked) TCP worker is created

src/main.c

@@ -103,8 +103,10 @@ int main (int argc, char* argv[])
 	// Save new queries to database (if database is used)
 	if(config.DBexport)
 	{
+		lock_shm();
 		if(DB_save_queries(NULL))
 			logg("Finished final database update");
+		unlock_shm();
 	}
 
 	cleanup(exit_code);