[Fix] Add configurable number of threads for OpenBLAS

Issue: #3082

5 files changed, 25 insertions, 2 deletions

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 6411601fa..797ae2074 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -206,6 +206,10 @@ IF (ENABLE_HYPERSCAN MATCHES "ON")
 	TARGET_LINK_LIBRARIES(rspamd-server hs)
 ENDIF()
 
+IF (WITH_BLAS)
+	TARGET_LINK_LIBRARIES(rspamd-server ${BLAS_REQUIRED_LIBRARIES})
+ENDIF()
+
 TARGET_LINK_LIBRARIES(rspamd-server ${RSPAMD_REQUIRED_LIBRARIES})
 
 ADD_EXECUTABLE(rspamd ${RSPAMDSRC} ${CMAKE_CURRENT_BINARY_DIR}/workers.c)
diff --git a/src/libserver/cfg_file.h b/src/libserver/cfg_file.h
index 4eea4db16..3ee832d5a 100644
--- a/src/libserver/cfg_file.h
+++ b/src/libserver/cfg_file.h
@@ -463,6 +463,7 @@ struct rspamd_config {
 	guint lua_gc_pause;                                /**< lua gc pause										*/
 	guint full_gc_iters;                            /**< iterations between full gc cycle					*/
 	guint max_lua_urls;                             /**< maximum number of urls to be passed to Lua			*/
+	guint max_blas_threads;                         /**< maximum threads for openblas when learning ANN		*/
 
 	GList *classify_headers;                        /**< list of headers using for statistics				*/
 	struct module_s **compiled_modules;                /**< list of compiled C modules							*/
diff --git a/src/libserver/cfg_rcl.c b/src/libserver/cfg_rcl.c
index 5ef913f36..7936c8725 100644
--- a/src/libserver/cfg_rcl.c
+++ b/src/libserver/cfg_rcl.c
@@ -2200,7 +2200,13 @@ rspamd_rcl_config_init (struct rspamd_config *cfg, GHashTable *skip_sections)
 				rspamd_rcl_parse_struct_integer,
 				G_STRUCT_OFFSET (struct rspamd_config, max_lua_urls),
 				RSPAMD_CL_FLAG_INT_32,
-				"Maximum count of URLs to pass to Lua to avoid DoS");
+				"Maximum count of URLs to pass to Lua to avoid DoS (default: 1024)");
+		rspamd_rcl_add_default_handler (sub,
+				"max_blas_threads",
+				rspamd_rcl_parse_struct_integer,
+				G_STRUCT_OFFSET (struct rspamd_config, max_blas_threads),
+				RSPAMD_CL_FLAG_INT_32,
+				"Maximum number of Blas threads for learning neural networks (default: 1)");
 
 		/* Neighbours configuration */
 		rspamd_rcl_add_section_doc (&sub->subsections, "neighbours", "name",
diff --git a/src/libserver/cfg_utils.c b/src/libserver/cfg_utils.c
index 5851e250f..98a50bea3 100644
--- a/src/libserver/cfg_utils.c
+++ b/src/libserver/cfg_utils.c
@@ -196,6 +196,7 @@ rspamd_config_new (enum rspamd_config_init_flags flags)
 	cfg->log_error_elt_maxlen = 1000;
 	cfg->cache_reload_time = 30.0;
 	cfg->max_lua_urls = 1024;
+	cfg->max_blas_threads = 1;
 
 	/* Default log line */
 	cfg->log_format_str = "id: <$mid>,$if_qid{ qid: <$>,}$if_ip{ ip: $,}"
diff --git a/src/libutil/util.c b/src/libutil/util.c
index ecdd7b2b9..5baffebd9 100644
--- a/src/libutil/util.c
+++ b/src/libutil/util.c
@@ -2401,6 +2401,14 @@ rspamd_free_zstd_dictionary (struct zstd_dictionary *dict)
 	}
 }
 
+#ifdef HAVE_CBLAS
+#ifdef HAVE_CBLAS_H
+#include "cblas.h"
+#else
+extern void openblas_set_num_threads(int num_threads);
+#endif
+#endif
+
 void
 rspamd_config_libs (struct rspamd_external_libs_ctx *ctx,
 		struct rspamd_config *cfg)
@@ -2494,6 +2502,9 @@ rspamd_config_libs (struct rspamd_external_libs_ctx *ctx,
 			ZSTD_freeCStream (ctx->out_zstream);
 			ctx->out_zstream = NULL;
 		}
+#ifdef HAVE_CBLAS
+		openblas_set_num_threads (cfg->max_blas_threads);
+#endif
 	}
 }
 
@@ -3267,4 +3278,4 @@ rspamd_set_counter_ema (struct rspamd_counter_data *cd,
 	cd->number ++;
 
 	return cd->mean;
-}
+}
\ No newline at end of file