Customizable system resources configuration

This change would allow end users to specify custom system resources configuration.
It should not change the existing behaviour as it keeps the defaults in place.

Signed-off-by: Maxime VISONNEAU <[email protected]>

Author: mvisonneau

consumer_config.go

@@ -245,7 +245,7 @@ func (r *configRoulette) genConfigs(bestCfg *consumerConfig, queueEmpty bool) {
 	r.oldBestCfg = bestCfg.Clone()
 	r.addConfig(r.oldBestCfg)
 
-	if !hasFreeSystemResources() {
+	if !hasFreeSystemResources(r.opt.MinSystemResources) {
 		internal.Logger.Println("taskq: system does not have enough free resources")
 		return
 	}

go.mod

@@ -19,6 +19,7 @@ require (
 	github.com/onsi/gomega v1.10.3
 	github.com/satori/go.uuid v1.2.0
 	github.com/vmihailenco/msgpack/v5 v5.0.0
+	github.com/stretchr/testify v1.6.1
 	golang.org/x/net v0.0.0-20201027133719-8eef5233e2a1 // indirect
 	google.golang.org/protobuf v1.25.0 // indirect
 )

queue.go

@@ -59,6 +59,9 @@ type QueueOptions struct {
 	// Optional message handler. The default is the global Tasks registry.
 	Handler Handler
 
+	// Minimal system resources required to consider the consumer available to process the queue
+	MinSystemResources SystemResources
+
 	inited bool
 }
 
@@ -118,6 +121,10 @@ func (opt *QueueOptions) Init() {
 	if opt.Handler == nil {
 		opt.Handler = &Tasks
 	}
+
+	if opt.MinSystemResources == (SystemResources{}) {
+		opt.MinSystemResources = NewDefaultSystemResources()
+	}
 }
 
 //------------------------------------------------------------------------------

sysinfo_linux.go

@@ -0,0 +1,28 @@
+package taskq
+
+const (
+	defaultSystemResourcesLoad1PerCPU          float64 = 1.5
+	defaultSystemResourcesMemoryFreeMB         uint64  = 2e5
+	defaultSystemResourcesMemoryFreePercentage uint64  = 5
+)
+
+// SystemResources represents system related values
+type SystemResources struct {
+	// Maximum per CPU load at 1min intervals
+	Load1PerCPU float64
+
+	// Minimum free memory required in megabytes
+	MemoryFreeMB uint64
+
+	// Minimum free memory required in percentage
+	MemoryFreePercentage uint64
+}
+
+// NewDefaultSystemResources returns a new SystemResources struct with some default values
+func NewDefaultSystemResources() SystemResources {
+	return SystemResources{
+		Load1PerCPU:          defaultSystemResourcesLoad1PerCPU,
+		MemoryFreeMB:         defaultSystemResourcesMemoryFreeMB,
+		MemoryFreePercentage: defaultSystemResourcesMemoryFreePercentage,
+	}
+}

sysinfo_other.go

@@ -0,0 +1,32 @@
+// +build linux
+
+package taskq
+
+import (
+	"runtime"
+
+	"github.com/capnm/sysinfo"
+	"github.com/vmihailenco/taskq/v3/internal"
+)
+
+func hasFreeSystemResources(sr SystemResources) bool {
+	si := sysinfo.Get()
+	free := si.FreeRam + si.BufferRam
+
+	if sr.Load1PerCPU > 0 && si.Loads[0] > sr.Load1PerCPU*float64(runtime.NumCPU()) {
+		internal.Logger.Println("taskq: consumer memory is lower than required")
+		return false
+	}
+
+	if sr.MemoryFreeMB > 0 && free < sr.MemoryFreeMB {
+		internal.Logger.Println("taskq: consumer memory is lower than required")
+		return false
+	}
+
+	if sr.MemoryFreePercentage > 0 && free/si.TotalRam < sr.MemoryFreePercentage/100 {
+		internal.Logger.Println("taskq: consumer memory is lower than required")
+		return false
+	}
+
+	return true
+}

system.go

@@ -0,0 +1,12 @@
+package taskq
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHasFreeSystemResources(t *testing.T) {
+	// TODO: Manage to mock capnm/sysinfo
+	assert.True(t, hasFreeSystemResources(SystemResources{}))
+}

system_linux.go

@@ -0,0 +1,7 @@
+// +build !linux
+
+package taskq
+
+func hasFreeSystemResources(_ SystemResources) bool {
+	return true
+}

system_linux_test.go

@@ -0,0 +1,11 @@
+package taskq
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHasFreeSystemResources(t *testing.T) {
+	assert.True(t, hasFreeSystemResources(SystemResources{}))
+}