Merge pull request #4052 from hashicorp/f-specify-total-memory

Allow to specify total memory on agent configuration

client/config/config.go

@@ -91,6 +91,10 @@ type Config struct {
 	// dynamically. It should be given as Cores * MHz (2 Cores * 2 Ghz = 4000)
 	CpuCompute int
 
+	// MemoryMB is the default node total memory in megabytes if it cannot be
+	// determined dynamically.
+	MemoryMB int
+
 	// MaxKillTimeout allows capping the user-specifiable KillTimeout. If the
 	// task's KillTimeout is greater than the MaxKillTimeout, MaxKillTimeout is
 	// used.

client/fingerprint/memory.go

@@ -9,6 +9,8 @@ import (
 	"github.com/shirou/gopsutil/mem"
 )
 
+const bytesInMB = 1024 * 1024
+
 // MemoryFingerprint is used to fingerprint the available memory on the node
 type MemoryFingerprint struct {
 	StaticFingerprinter
@@ -24,17 +26,25 @@ func NewMemoryFingerprint(logger *log.Logger) Fingerprint {
 }
 
 func (f *MemoryFingerprint) Fingerprint(req *cstructs.FingerprintRequest, resp *cstructs.FingerprintResponse) error {
-	memInfo, err := mem.VirtualMemory()
-	if err != nil {
-		f.logger.Printf("[WARN] Error reading memory information: %s", err)
-		return err
+	var totalMemory int
+	cfg := req.Config
+	if cfg.MemoryMB != 0 {
+		totalMemory = cfg.MemoryMB * bytesInMB
+	} else {
+		memInfo, err := mem.VirtualMemory()
+		if err != nil {
+			f.logger.Printf("[WARN] Error reading memory information: %s", err)
+			return err
+		}
+		if memInfo.Total > 0 {
+			totalMemory = int(memInfo.Total)
+		}
 	}
 
-	if memInfo.Total > 0 {
-		resp.AddAttribute("memory.totalbytes", fmt.Sprintf("%d", memInfo.Total))
-
+	if totalMemory > 0 {
+		resp.AddAttribute("memory.totalbytes", fmt.Sprintf("%d", totalMemory))
 		resp.Resources = &structs.Resources{
-			MemoryMB: int(memInfo.Total / 1024 / 1024),
+			MemoryMB: totalMemory / bytesInMB,
 		}
 	}
 

client/fingerprint/memory_test.go

@@ -6,6 +6,8 @@ import (
 	"github.com/hashicorp/nomad/client/config"
 	cstructs "github.com/hashicorp/nomad/client/structs"
 	"github.com/hashicorp/nomad/nomad/structs"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestMemoryFingerprint(t *testing.T) {
@@ -30,3 +32,23 @@ func TestMemoryFingerprint(t *testing.T) {
 		t.Fatalf("Expected node.Resources.MemoryMB to be non-zero")
 	}
 }
+
+func TestMemoryFingerprint_Override(t *testing.T) {
+	f := NewMemoryFingerprint(testLogger())
+	node := &structs.Node{
+		Attributes: make(map[string]string),
+	}
+
+	memoryMB := 15000
+	request := &cstructs.FingerprintRequest{Config: &config.Config{MemoryMB: memoryMB}, Node: node}
+	var response cstructs.FingerprintResponse
+	err := f.Fingerprint(request, &response)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	assertNodeAttributeContains(t, response.Attributes, "memory.totalbytes")
+	require := require.New(t)
+	require.NotNil(response.Resources)
+	require.Equal(response.Resources.MemoryMB, memoryMB)
+}

command/agent/agent.go

@@ -356,6 +356,9 @@ func (a *Agent) clientConfig() (*clientconfig.Config, error) {
 	if a.config.Client.CpuCompute != 0 {
 		conf.CpuCompute = a.config.Client.CpuCompute
 	}
+	if a.config.Client.MemoryMB != 0 {
+		conf.MemoryMB = a.config.Client.MemoryMB
+	}
 	if a.config.Client.MaxKillTimeout != "" {
 		dur, err := time.ParseDuration(a.config.Client.MaxKillTimeout)
 		if err != nil {

command/agent/config-test-fixtures/non-optional.hcl

@@ -0,0 +1,3 @@
+client {
+	memory_total_mb = 5555
+}

command/agent/config.go

@@ -175,6 +175,9 @@ type ClientConfig struct {
 	// CpuCompute is used to override any detected or default total CPU compute.
 	CpuCompute int `mapstructure:"cpu_total_compute"`
 
+	// MemoryMB is used to override any detected or default total memory.
+	MemoryMB int `mapstructure:"memory_total_mb"`
+
 	// MaxKillTimeout allows capping the user-specifiable KillTimeout.
 	MaxKillTimeout string `mapstructure:"max_kill_timeout"`
 
@@ -1094,6 +1097,9 @@ func (a *ClientConfig) Merge(b *ClientConfig) *ClientConfig {
 	if b.CpuCompute != 0 {
 		result.CpuCompute = b.CpuCompute
 	}
+	if b.MemoryMB != 0 {
+		result.MemoryMB = b.MemoryMB
+	}
 	if b.MaxKillTimeout != "" {
 		result.MaxKillTimeout = b.MaxKillTimeout
 	}

command/agent/config_parse.go

@@ -356,6 +356,7 @@ func parseClient(result **ClientConfig, list *ast.ObjectList) error {
 		"chroot_env",
 		"network_interface",
 		"network_speed",
+		"memory_total_mb",
 		"cpu_total_compute",
 		"max_kill_timeout",
 		"client_max_port",

command/agent/config_parse_test.go

@@ -2,14 +2,12 @@ package agent
 
 import (
 	"path/filepath"
-	"reflect"
-	"strings"
 	"testing"
 	"time"
 
 	"github.com/hashicorp/nomad/helper"
 	"github.com/hashicorp/nomad/nomad/structs/config"
-	"github.com/kr/pretty"
+	"github.com/stretchr/testify/require"
 )
 
 func TestConfig_Parse(t *testing.T) {
@@ -64,6 +62,7 @@ func TestConfig_Parse(t *testing.T) {
 					NetworkInterface: "eth0",
 					NetworkSpeed:     100,
 					CpuCompute:       4444,
+					MemoryMB:         0,
 					MaxKillTimeout:   "10s",
 					ClientMinPort:    1000,
 					ClientMaxPort:    2000,
@@ -206,9 +205,64 @@ func TestConfig_Parse(t *testing.T) {
 			},
 			false,
 		},
+		{
+			"non-optional.hcl",
+			&Config{
+				Region:         "",
+				Datacenter:     "",
+				NodeName:       "",
+				DataDir:        "",
+				LogLevel:       "",
+				BindAddr:       "",
+				EnableDebug:    false,
+				Ports:          nil,
+				Addresses:      nil,
+				AdvertiseAddrs: nil,
+				Client: &ClientConfig{
+					Enabled:               false,
+					StateDir:              "",
+					AllocDir:              "",
+					Servers:               nil,
+					NodeClass:             "",
+					Meta:                  nil,
+					Options:               nil,
+					ChrootEnv:             nil,
+					NetworkInterface:      "",
+					NetworkSpeed:          0,
+					CpuCompute:            0,
+					MemoryMB:              5555,
+					MaxKillTimeout:        "",
+					ClientMinPort:         0,
+					ClientMaxPort:         0,
+					Reserved:              nil,
+					GCInterval:            0,
+					GCParallelDestroys:    0,
+					GCDiskUsageThreshold:  0,
+					GCInodeUsageThreshold: 0,
+					GCMaxAllocs:           0,
+					NoHostUUID:            nil,
+				},
+				Server:                    nil,
+				ACL:                       nil,
+				Telemetry:                 nil,
+				LeaveOnInt:                false,
+				LeaveOnTerm:               false,
+				EnableSyslog:              false,
+				SyslogFacility:            "",
+				DisableUpdateCheck:        nil,
+				DisableAnonymousSignature: false,
+				Consul:                 nil,
+				Vault:                  nil,
+				TLSConfig:              nil,
+				HTTPAPIResponseHeaders: nil,
+				Sentinel:               nil,
+			},
+			false,
+		},
 	}
 
 	for _, tc := range cases {
+		require := require.New(t)
 		t.Run(tc.File, func(t *testing.T) {
 			path, err := filepath.Abs(filepath.Join("./config-test-fixtures", tc.File))
 			if err != nil {
@@ -219,10 +273,7 @@ func TestConfig_Parse(t *testing.T) {
 			if (err != nil) != tc.Err {
 				t.Fatalf("file: %s\n\n%s", tc.File, err)
 			}
-
-			if !reflect.DeepEqual(actual, tc.Result) {
-				t.Errorf("file: %s  diff: (actual vs expected)\n\n%s", tc.File, strings.Join(pretty.Diff(actual, tc.Result), "\n"))
-			}
+			require.EqualValues(actual, tc.Result)
 		})
 	}
 }

command/agent/config_test.go

@@ -86,6 +86,7 @@ func TestConfig_Merge(t *testing.T) {
 			},
 			NetworkSpeed:   100,
 			CpuCompute:     100,
+			MemoryMB:       100,
 			MaxKillTimeout: "20s",
 			ClientMaxPort:  19996,
 			Reserved: &Resources{
@@ -232,6 +233,7 @@ func TestConfig_Merge(t *testing.T) {
 			ClientMinPort:  22000,
 			NetworkSpeed:   105,
 			CpuCompute:     105,
+			MemoryMB:       105,
 			MaxKillTimeout: "50s",
 			Reserved: &Resources{
 				CPU:                 15,

website/source/docs/agent/configuration/client.html.md

@@ -67,6 +67,9 @@ client {
   clients can determine their total CPU compute automatically, and thus in most
   cases this should be left unset.
 
+- `memory_total_mb` `(int:0)` - Specifies an override for the total memory. If set,
+  this value overrides any detected memory.
+
 - `node_class` `(string: "")` - Specifies an arbitrary string used to logically
   group client nodes by user-defined class. This can be used during job
   placement as a filter.