From fafb320a3f5d93876b9b35e4575a0e85924c0817 Mon Sep 17 00:00:00 2001 From: Kent Quirk Date: Thu, 26 Oct 2023 15:51:54 -0400 Subject: [PATCH] fix: Fix AvailableMemory parsing on the cmd line (#875) ## Which problem is this PR solving? - #857 ## Short description of the changes - Change AvailableMemory in cmdenv to be type MemorySize - Implement flags.Unmarshaler so the flags system uses it - Add a test to make sure it works Fixes #857 --- config/cmdenv.go | 46 +++++++++++++++++++++---------------------- config/config_test.go | 15 ++++++++++++++ config/file_config.go | 8 ++++++++ 3 files changed, 46 insertions(+), 23 deletions(-) diff --git a/config/cmdenv.go b/config/cmdenv.go index 5f70edbb09..df45f32ec8 100644 --- a/config/cmdenv.go +++ b/config/cmdenv.go @@ -26,29 +26,29 @@ import ( // that this system uses reflection to establish the relationship between the // config struct and the command line options. type CmdEnv struct { - ConfigLocation string `short:"c" long:"config" env:"REFINERY_CONFIG" default:"/etc/refinery/refinery.yaml" description:"config file or URL to load"` - RulesLocation string `short:"r" long:"rules_config" env:"REFINERY_RULES_CONFIG" default:"/etc/refinery/rules.yaml" description:"config file or URL to load"` - HTTPListenAddr string `long:"http-listen-address" env:"REFINERY_HTTP_LISTEN_ADDRESS" description:"HTTP listen address for incoming event traffic"` - PeerListenAddr string `long:"peer-listen-address" env:"REFINERY_PEER_LISTEN_ADDRESS" description:"Peer listen address for communication between Refinery instances"` - GRPCListenAddr string `long:"grpc-listen-address" env:"REFINERY_GRPC_LISTEN_ADDRESS" description:"gRPC listen address for OTLP traffic"` - RedisHost string `long:"redis-host" env:"REFINERY_REDIS_HOST" description:"Redis host address"` - RedisUsername string `long:"redis-username" env:"REFINERY_REDIS_USERNAME" description:"Redis username"` - RedisPassword string `long:"redis-password" env:"REFINERY_REDIS_PASSWORD" description:"Redis password"` - RedisAuthCode string `long:"redis-auth-code" env:"REFINERY_REDIS_AUTH_CODE" description:"Redis AUTH code"` - HoneycombAPI string `long:"honeycomb-api" env:"REFINERY_HONEYCOMB_API" description:"Honeycomb API URL"` - HoneycombAPIKey string `long:"honeycomb-api-key" env:"REFINERY_HONEYCOMB_API_KEY" description:"Honeycomb API key (for logger and metrics)"` - HoneycombLoggerAPIKey string `long:"logger-api-key" env:"REFINERY_HONEYCOMB_LOGGER_API_KEY" description:"Honeycomb logger API key"` - LegacyMetricsAPIKey string `long:"legacy-metrics-api-key" env:"REFINERY_HONEYCOMB_METRICS_API_KEY" description:"API key for legacy Honeycomb metrics"` - OTelMetricsAPIKey string `long:"otel-metrics-api-key" env:"REFINERY_OTEL_METRICS_API_KEY" description:"API key for OTel metrics if being sent to Honeycomb"` - QueryAuthToken string `long:"query-auth-token" env:"REFINERY_QUERY_AUTH_TOKEN" description:"Token for debug/management queries"` - AvailableMemory string `long:"available-memory" env:"REFINERY_AVAILABLE_MEMORY" description:"The maximum memory available for Refinery to use."` - Debug bool `short:"d" long:"debug" description:"Runs debug service (on the first open port between localhost:6060 and :6069 by default)"` - Version bool `short:"v" long:"version" description:"Print version number and exit"` - InterfaceNames bool `long:"interface-names" description:"Print system's network interface names and exit."` - Validate bool `short:"V" long:"validate" description:"Validate the configuration files, writing results to stdout, and exit with 0 if valid, 1 if invalid."` - NoValidate bool `long:"no-validate" description:"Do not attempt to validate the configuration files. Makes --validate meaningless."` - WriteConfig string `long:"write-config" description:"After applying defaults, environment variables, and command line values, write the loaded configuration to the specified file as YAML and exit."` - WriteRules string `long:"write-rules" description:"After applying defaults, write the loaded rules to the specified file as YAML and exit."` + ConfigLocation string `short:"c" long:"config" env:"REFINERY_CONFIG" default:"/etc/refinery/refinery.yaml" description:"config file or URL to load"` + RulesLocation string `short:"r" long:"rules_config" env:"REFINERY_RULES_CONFIG" default:"/etc/refinery/rules.yaml" description:"config file or URL to load"` + HTTPListenAddr string `long:"http-listen-address" env:"REFINERY_HTTP_LISTEN_ADDRESS" description:"HTTP listen address for incoming event traffic"` + PeerListenAddr string `long:"peer-listen-address" env:"REFINERY_PEER_LISTEN_ADDRESS" description:"Peer listen address for communication between Refinery instances"` + GRPCListenAddr string `long:"grpc-listen-address" env:"REFINERY_GRPC_LISTEN_ADDRESS" description:"gRPC listen address for OTLP traffic"` + RedisHost string `long:"redis-host" env:"REFINERY_REDIS_HOST" description:"Redis host address"` + RedisUsername string `long:"redis-username" env:"REFINERY_REDIS_USERNAME" description:"Redis username"` + RedisPassword string `long:"redis-password" env:"REFINERY_REDIS_PASSWORD" description:"Redis password"` + RedisAuthCode string `long:"redis-auth-code" env:"REFINERY_REDIS_AUTH_CODE" description:"Redis AUTH code"` + HoneycombAPI string `long:"honeycomb-api" env:"REFINERY_HONEYCOMB_API" description:"Honeycomb API URL"` + HoneycombAPIKey string `long:"honeycomb-api-key" env:"REFINERY_HONEYCOMB_API_KEY" description:"Honeycomb API key (for logger and metrics)"` + HoneycombLoggerAPIKey string `long:"logger-api-key" env:"REFINERY_HONEYCOMB_LOGGER_API_KEY" description:"Honeycomb logger API key"` + LegacyMetricsAPIKey string `long:"legacy-metrics-api-key" env:"REFINERY_HONEYCOMB_METRICS_API_KEY" description:"API key for legacy Honeycomb metrics"` + OTelMetricsAPIKey string `long:"otel-metrics-api-key" env:"REFINERY_OTEL_METRICS_API_KEY" description:"API key for OTel metrics if being sent to Honeycomb"` + QueryAuthToken string `long:"query-auth-token" env:"REFINERY_QUERY_AUTH_TOKEN" description:"Token for debug/management queries"` + AvailableMemory MemorySize `long:"available-memory" env:"REFINERY_AVAILABLE_MEMORY" description:"The maximum memory available for Refinery to use (ex: 4GiB)."` + Debug bool `short:"d" long:"debug" description:"Runs debug service (on the first open port between localhost:6060 and :6069 by default)"` + Version bool `short:"v" long:"version" description:"Print version number and exit"` + InterfaceNames bool `long:"interface-names" description:"Print system's network interface names and exit."` + Validate bool `short:"V" long:"validate" description:"Validate the configuration files, writing results to stdout, and exit with 0 if valid, 1 if invalid."` + NoValidate bool `long:"no-validate" description:"Do not attempt to validate the configuration files. Makes --validate meaningless."` + WriteConfig string `long:"write-config" description:"After applying defaults, environment variables, and command line values, write the loaded configuration to the specified file as YAML and exit."` + WriteRules string `long:"write-rules" description:"After applying defaults, write the loaded rules to the specified file as YAML and exit."` } func NewCmdEnvOptions(args []string) (*CmdEnv, error) { diff --git a/config/config_test.go b/config/config_test.go index c97f72315e..7b2284397a 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -452,6 +452,21 @@ func TestMaxAlloc(t *testing.T) { assert.Equal(t, expected, inMemConfig.MaxAlloc) } +func TestAvailableMemoryCmdLine(t *testing.T) { + cm := makeYAML("General.ConfigurationVersion", 2, "Collection.CacheCapacity", 1000, "Collection.AvailableMemory", 2_000_000_000) + rm := makeYAML("ConfigVersion", 2) + config, rules := createTempConfigs(t, cm, rm) + defer os.Remove(rules) + defer os.Remove(config) + c, err := getConfig([]string{"--no-validate", "--config", config, "--rules_config", rules, "--available-memory", "8Gib"}) + assert.NoError(t, err) + + expected := MemorySize(8 * 1024 * 1024 * 1024) + inMemConfig, err := c.GetCollectionConfig() + assert.NoError(t, err) + assert.Equal(t, expected, inMemConfig.AvailableMemory) +} + func TestGetSamplerTypes(t *testing.T) { cm := makeYAML("General.ConfigurationVersion", 2) rm := makeYAML( diff --git a/config/file_config.go b/config/file_config.go index 4b62085490..d39fddd9d3 100644 --- a/config/file_config.go +++ b/config/file_config.go @@ -13,6 +13,7 @@ import ( "sync" "time" + "github.com/jessevdk/go-flags" "gopkg.in/yaml.v3" ) @@ -175,6 +176,13 @@ func (m *MemorySize) UnmarshalText(text []byte) error { return nil } +// Make sure we implement flags.Unmarshaler so that it works with cmdenv. +var _ flags.Unmarshaler = (*MemorySize)(nil) + +func (m *MemorySize) UnmarshalFlag(value string) error { + return m.UnmarshalText([]byte(value)) +} + type fileConfig struct { mainConfig *configContents mainHash string