Customizable system resources configuration

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()
+	}
 }
 
 //------------------------------------------------------------------------------

system.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,
+	}
+}

system_linux.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_linux_test.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_other.go

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

system_other_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{}))
+}

system_test.go

@@ -0,0 +1,16 @@
+package taskq
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewDefaultSystemResources(t *testing.T) {
+	expectedValue := SystemResources{
+		Load1PerCPU:          1.5,
+		MemoryFreeMB:         2e5,
+		MemoryFreePercentage: 5,
+	}
+	assert.Equal(t, expectedValue, NewDefaultSystemResources())
+}