From 28ecc74d7a29077148f426842e2ddf5220f90feb Mon Sep 17 00:00:00 2001
From: Sacha Ligthert <sacha@ligthert.net>
Date: Sat, 25 Jan 2025 19:45:31 +0100
Subject: [PATCH] Allows setting of cores to use (sorta fixes #8)

---
 solver/flags.go  | 12 ++++++++++++
 solver/solver.go |  7 +++++++
 solver/types.go  |  5 ++++-
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/solver/flags.go b/solver/flags.go
index 081d029..e5d0ce0 100644
--- a/solver/flags.go
+++ b/solver/flags.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"runtime"
 )
 
 func (solver *Solver) parse_flags() {
@@ -30,10 +31,21 @@ func (solver *Solver) parse_flags() {
 	flag.StringVar(&row7, "row7", "000000000", "7th row of the sudoku puzzle.")
 	flag.StringVar(&row8, "row8", "000000000", "8th row of the sudoku puzzle.")
 	flag.StringVar(&row9, "row9", "000000000", "9th row of the sudoku puzzle.")
+	flag.IntVar(&solver.numcpus, "numcpu", runtime.NumCPU(), "Number of CPU cores to assign to this task.")
 
 	// Parse the flags
 	flag.Parse()
 
+	if solver.numcpus <= 0 {
+		log.Printf("ERROR: Number of CPU cores must be 1 or higher.\n\n")
+		solver.print_Usage()
+		os.Exit(1)
+	}
+
+	if solver.numcpus != runtime.NumCPU() {
+		runtime.GOMAXPROCS(solver.numcpus)
+	}
+
 	if row1 == "000000000" || row2 == "000000000" || row3 == "000000000" || row4 == "000000000" || row5 == "000000000" || row6 == "000000000" || row7 == "000000000" || row8 == "000000000" || row9 == "000000000" {
 		log.Printf("ERROR: All parameters must be entered.\n\n")
 		solver.print_Usage()
diff --git a/solver/solver.go b/solver/solver.go
index 3bba782..bd05c94 100644
--- a/solver/solver.go
+++ b/solver/solver.go
@@ -2,6 +2,8 @@ package solver
 
 import (
 	"log"
+	"runtime"
+	"strconv"
 )
 
 func Run() {
@@ -11,6 +13,11 @@ func Run() {
 	// Parse and handle flags
 	solver.parse_flags()
 
+	// Report number of CPUs being used, if set.
+	if runtime.NumCPU() != solver.numcpus {
+		log.Println("Using " + strconv.Itoa(solver.numcpus) + " CPUs, (was " + strconv.Itoa(runtime.NumCPU()) + ")")
+	}
+
 	// Load blocks from CSV file
 	solver.load_blocks()
 
diff --git a/solver/types.go b/solver/types.go
index 649efa9..a9ffa00 100644
--- a/solver/types.go
+++ b/solver/types.go
@@ -1,6 +1,8 @@
 package solver
 
-import "sync/atomic"
+import (
+	"sync/atomic"
+)
 
 type Solver struct {
 	blocks    []int
@@ -26,4 +28,5 @@ type Solver struct {
 	counter   atomic.Int64
 	solutions []string
 	rates     []int64
+	numcpus   int
 }