From 9ddccba7807172afbed5297932c7522daf15ce94 Mon Sep 17 00:00:00 2001
From: m0duspwnens <josh.patterson@securityonionsolutions.com>
Date: Wed, 28 Aug 2024 10:09:42 -0400
Subject: [PATCH] LSHEAP and pipeline workers for virt

---
 salt/manager/tools/sbin/so-minion   |  6 +++++-
 salt/reactor/sominion_setup.sls     |  2 +-
 salt/setup/virt/soinstall.map.jinja | 12 +++++++++++-
 salt/setup/virt/sominion.sls        |  2 +-
 4 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/salt/manager/tools/sbin/so-minion b/salt/manager/tools/sbin/so-minion
index 18c8c3224..0956840b1 100755
--- a/salt/manager/tools/sbin/so-minion
+++ b/salt/manager/tools/sbin/so-minion
@@ -44,6 +44,10 @@ for i in "$@"; do
 			ES_HEAP_SIZE="${i#*=}"
 			shift
 			;;
+		-l=*|--lsheap=*)
+			LSHEAP="${i#*=}"
+			shift
+			;;
 		-n=*|--mgmtnic=*)
 			MNIC="${i#*=}"
 			shift
@@ -214,7 +218,7 @@ function add_logstash_to_minion() {
 		"logstash:"\
 		"  enabled: True"\
 		"  config:"\
-		"    pipeline_x_workers: $CPUCORES"\
+		"    pipeline_x_workers: $CORECOUNT"\
 		"  settings:"\
 		"    lsheap: $LSHEAP"\
 		"  " >> $PILLARFILE
diff --git a/salt/reactor/sominion_setup.sls b/salt/reactor/sominion_setup.sls
index 91a4a6cb9..a78475e6c 100644
--- a/salt/reactor/sominion_setup.sls
+++ b/salt/reactor/sominion_setup.sls
@@ -28,7 +28,7 @@ def run():
   with open("/opt/so/saltstack/local/pillar/hypervisor/" + hv_name + "/" + minionid + ".sls", 'w') as f:
     yaml.dump(vm_out_data, f, default_flow_style=False)
 
-  rc = call("NODETYPE=" + DATA['NODETYPE'] + " /usr/sbin/so-minion -o=addVirt -m=" + minionid + " -n=" + DATA['MNIC'] + " -i=" + DATA['MAINIP'] + " -a=" + DATA['INTERFACE'] + " -c=" + str(DATA['CORECOUNT'])  + " -d='" + DATA['NODE_DESCRIPTION'] + "'", shell=True)
+  rc = call("NODETYPE=" + DATA['NODETYPE'] + " /usr/sbin/so-minion -o=addVirt -m=" + minionid + " -n=" + DATA['MNIC'] + " -i=" + DATA['MAINIP'] + " -a=" + DATA['INTERFACE'] + " -c=" + str(DATA['CPU'])  + " -d='" + DATA['NODE_DESCRIPTION'] + "'" + " -e=" + DATA['ES_HEAP_SIZE'] + " -l=" + DATA['LS_HEAP_SIZE'], shell=True)
 
   logging.error('sominion_setup reactor: rc: %s' % rc)
 
diff --git a/salt/setup/virt/soinstall.map.jinja b/salt/setup/virt/soinstall.map.jinja
index e2649aab1..a589250df 100644
--- a/salt/setup/virt/soinstall.map.jinja
+++ b/salt/setup/virt/soinstall.map.jinja
@@ -14,7 +14,8 @@
 
 {% if nodetype == "searchnode" %}
 
-{%   do DATA.update({'LSHOSTNAME': grains.host}) %}
+{#   we can't use the host grain here because the grain may not be updated yet from the hostname change #}
+{%   do DATA.update({'LSHOSTNAME': salt['cmd.run']('hostname')}) %}
 
 {#   this replicates the function es_heapsize in so-functions #}
 {%   if total_mem < 8000 %}
@@ -31,4 +32,13 @@
 {%   endif %}
 {%   do DATA.update({'ES_HEAP_SIZE': ES_HEAP_SIZE}) %}
 
+{%   if total_mem >= 32000 or nodetype in ['managersearch','heavynode','standalone'] %}
+{%     set LS_HEAP_SIZE="1000m" %}
+{%   elif nodetype == 'eval' %}
+{%     set LS_HEAP_SIZE="700m" %}
+{%   else %}
+{%     set LS_HEAP_SIZE="500m" %}
+{%   endif %}
+{%   do DATA.update({'LS_HEAP_SIZE': LS_HEAP_SIZE}) %}
+
 {% endif %}
diff --git a/salt/setup/virt/sominion.sls b/salt/setup/virt/sominion.sls
index 328f50311..55e3a796a 100644
--- a/salt/setup/virt/sominion.sls
+++ b/salt/setup/virt/sominion.sls
@@ -19,7 +19,7 @@ create_pillar:
         NODETYPE: {{ DATA.NODETYPE }}
         CORECOUNT: {{ DATA.CORECOUNT }}
         LSHOSTNAME: {{ DATA.LSHOSTNAME }}
-        LSHEAP: {{ DATA.LSHEAP }}
+        LS_HEAP_SIZE: {{ DATA.LSHEAP }}
         CPUCORES: {{ DATA.CPUCORES }}
         IDH_MGTRESTRICT: {{ DATA.IDH_MGTRESTRICT }}
         IDH_SERVICES: {{ DATA.IDH_SERVICES }}