Repository: honeycombio/refinery
Pull request #721 (2023-06-17 20:29:42)

feat: MaxAlloc improvements
Dependent on #719; can be reviewed before merging that.

## Which problem is this PR solving?

In v1, Refinery allowed setting a `MaxAlloc` value that controlled the maximum amount of memory it would try to use. This is hard to keep configured properly, because the guidance is to set it to some fraction of the memory allocated to the pod. This PR implements two new values: 

- the full allocation amount, called `AvailableMemory`
- the percentage refinery should use, called `MaxMemory`

The existing `MaxAlloc` value is still supported. If the first two are nonzero, they are multiplied to get the maximum allocation; otherwise, `MaxAlloc` is used. As in v1, if all of them are 0, refinery doesn't attempt to limit its own use.

I spent a while trying to find a way to allow the app to read the max allocation from the system, but it's pretty platform-dependent. It seemed better to just make it so that the value can be put into the environment in the Helm chart.

## Short description of the changes

- Add `AvailableMemory` and make it settable in the environment and on the command line
- Implement the logic for calculating the value refinery should use and use that in the allocation check
- Fix the names of the appropriate function calls and mocks
- Regenerate files



Patch:
diff --git a/collect/collect.go b/collect/collect.go
index 228229c..1ea23af 100644
--- a/collect/collect.go
+++ b/collect/collect.go
@@ -74,7 +74,7 @@ type InMemCollector struct {
 func (i *InMemCollector) Start() error {
 	i.Logger.Debug().Logf("Starting InMemCollector")
 	defer func() { i.Logger.Debug().Logf("Finished starting InMemCollector") }()
-	imcConfig, err := i.Config.GetInMemCollectorCacheCapacity()
+	imcConfig, err := i.Config.GetCollectionConfig()
 	if err != nil {
 		return err
 	}
@@ -149,7 +149,7 @@ func (i *InMemCollector) sendReloadSignal() {
 
 func (i *InMemCollector) reloadConfigs() {
 	i.Logger.Debug().Logf("reloading in-mem collect config")
-	imcConfig, err := i.Config.GetInMemCollectorCacheCapacity()
+	imcConfig, err := i.Config.GetCollectionConfig()
 	if err != nil {
 		i.Logger.Error().WithField("error", err).Logf("Failed to reload InMemCollector section when reloading configs")
 	}
@@ -185,13 +185,14 @@ func (i *InMemCollector) reloadConfigs() {
 }
 
 func (i *InMemCollector) checkAlloc() {
-	inMemConfig, err := i.Config.GetInMemCollectorCacheCapacity()
-	i.Metrics.Store("MEMORY_MAX_ALLOC", float64(inMemConfig.MaxAlloc))
+	inMemConfig, err := i.Config.GetCollectionConfig()
+	maxAlloc := inMemConfig.GetMaxAlloc()
+	i.Metrics.Store("MEMORY_MAX_ALLOC", float64(maxAlloc))
 
 	var mem runtime.MemStats
 	runtime.ReadMemStats(&mem)
 	i.Metrics.Gauge("memory_heap_allocation", int64(mem.Alloc))
-	if err != nil || inMemConfig.MaxAlloc == 0 || mem.Alloc < uint64(inMemConfig.MaxAlloc) {
+	if err != nil || maxAlloc == 0 || mem.Alloc < uint64(maxAlloc) {
 		return
 	}
 
@@ -199,7 +200,7 @@ func (i *InMemCollector) checkAlloc() {
 	// enough to get us below the max capacity, but not TOO much below.
 	// Because our impact numbers are only the data size, reducing by enough to reach
 	// max alloc will actually do more than that.
-	totalToRemove := mem.Alloc - uint64(inMemConfig.MaxAlloc)
+	totalToRemove := mem.Alloc - uint64(maxAlloc)
 
 	// The size of the cache exceeds the user's intended allocation, so we're going to
 	// remove the traces from the cache that have had the most impact on allocation.
diff --git a/config.md b/config.md
index eaa6778..52a6fb9 100644
--- a/config.md
+++ b/config.md
@@ -1,7 +1,7 @@
 # Honeycomb Refinery Configuration Documentation
 
 This is the documentation for the configuration file for Honeycomb's Refinery.
-It was automatically generated on 2023-06-20 at 16:34:41 UTC.
+It was automatically generated on 2023-06-20 at 18:40:47 UTC.
 
 ## The Config file
 
@@ -778,6 +778,8 @@ Refinery will timeout after this duration when communicating with Redis.
 ### Section Name: `Collection`
 
 Brings together the settings that are relevant to collecting spans together to make traces.
+If none of the memory settings are used, then Refinery will not attempt to limit its memory usage.
+This is not recommended for production use since a burst of traffic could cause Refinery to run out of memory and crash.
 
 
 
@@ -792,18 +794,30 @@ The number of traces in the cache should be many multiples (100x to 1000x) of th
 - Type: `int`
 - Default: `10000`
 
-### `MaxMemory`
+### `AvailableMemory`
 
-MaxMemory is the maximum percentage of memory that should be allocated by the span collector.
+AvailableMemory is the amount of system memory available to the refinery process.
+
+The amount of system memory available to the refinery process.
+This value will typically be set through an environment variable controlled by the container or deploy script.
+If this value is zero or not set, MaxMemory cannot be used to calculate the maximum allocation and MaxAlloc will be used instead.
+If set, this must be a memory size.
+64-bit values are supported.
+Sizes with standard unit suffixes like "MB" and "GiB" are also supported.
+
+- Eligible for live reload.
+- Type: `memorysize`
+- Example: `4Gb`
+
+### `MaxMemoryPercentage`
+
+MaxMemoryPercentage is the maximum percentage of memory that should be allocated by the span collector.
 
 If nonzero, it must be an integer value between 1 and 100, representing the target maximum percentage of memory that should be allocated by the span collector.
 If set to a non-zero value, once per tick (see SendTicker) the collector will compare total allocated bytes to this calculated value.
 If allocation is too high, traces will be ejected from the cache early to reduce memory.
 Useful values for this setting are generally in the range of 70-90.
-Depending on deployment details, system memory information may not be available.
-If it is not, a warning will be logged and the value of MaxAlloc will be used.
 If this value is 0, MaxAlloc will be used.
-Requires MaxAlloc to be nonzero.
 
 - Eligible for live reload.
 - Type: `percentage`
diff --git a/config/cmdenv.go b/config/cmdenv.go
index 60c4c7a..7c88b45 100644
--- a/config/cmdenv.go
+++ b/config/cmdenv.go
@@ -40,6 +40,7 @@ type CmdEnv struct {
 	LegacyMetricsAPIKey   string `long:"legacy-metrics-api-key" env:"REFINERY_LEGACY_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."`
diff --git a/config/config.go b/config/config.go
index c0c700b..fca8ab6 100644
--- a/config/config.go
+++ b/config/config.go
@@ -102,8 +102,8 @@ type Config interface {
 	// GetHoneycombLoggerConfig returns the config specific to the HoneycombLogger
 	GetHoneycombLoggerConfig() (HoneycombLoggerConfig, error)
 
-	// GetInMemCollectorCacheCapacity returns the config specific to the InMemCollector
-	GetInMemCollectorCacheCapacity() (CollectionConfig, error)
+	// GetCollectionConfig returns the config specific to the InMemCollector
+	GetCollectionConfig() (CollectionConfig, error)
 
 	// GetSamplerConfigForDestName returns the sampler type and name to use for
 	// the given destination (environment, or dataset in classic)
diff --git a/config/file_config.go b/config/file_config.go
index 5c844a7..7330520 100644
--- a/config/file_config.go
+++ b/config/file_config.go
@@ -184,9 +184,17 @@ type RedisPeerManagementConfig struct {
 
 type CollectionConfig struct {
 	// CacheCapacity must be less than math.MaxInt32
-	CacheCapacity int        `yaml:"CacheCapacity" default:"10_000"`
-	MaxMemory     int        `yaml:"MaxMemory" default:"75"`
-	MaxAlloc      MemorySize `yaml:"MaxAlloc"`
+	CacheCapacity       int        `yaml:"CacheCapacity" default:"10_000"`
+	AvailableMemory     MemorySize `yaml:"AvailableMemory" cmdenv:"AvailableMemory"`
+	MaxMemoryPercentage int        `yaml:"MaxMemoryPercentage" default:"75"`
+	MaxAlloc            MemorySize `yaml:"MaxAlloc"`
+}
+
+func (c CollectionConfig) GetMaxAlloc() MemorySize {
+	if c.AvailableMemory == 0 || c.MaxMemoryPercentage == 0 {
+		return c.MaxAlloc
+	}
+	return c.AvailableMemory * MemorySize(c.MaxMemoryPercentage) / 100
 }
 
 type BufferSizeConfig struct {
@@ -621,7 +629,7 @@ func (f *fileConfig) GetSamplerConfigForDestName(destname string) (any, string,
 	return cfg, name, err
 }
 
-func (f *fileConfig) GetInMemCollectorCacheCapacity() (CollectionConfig, error) {
+func (f *fileConfig) GetCollectionConfig() (CollectionConfig, error) {
 	f.mux.RLock()
 	defer f.mux.RUnlock()
 
diff --git a/config/metadata/configMeta.yaml b/config/metadata/configMeta.yaml
index 5dd9a45..6c4f823 100644
--- a/config/metadata/configMeta.yaml
+++ b/config/metadata/configMeta.yaml
@@ -873,8 +873,11 @@ groups:
   - name: Collection
     title: "Collection Settings"
     description: >
-      Brings together the settings that are relevant to
-      collecting spans together to make traces.
+      Brings together the settings that are relevant to collecting spans
+      together to make traces. If none of the memory settings are used, then
+      Refinery will not attempt to limit its memory usage. This is not
+      recommended for production use since a burst of traffic could cause
+      Refinery to run out of memory and crash.
     fields:
       - name: CacheCapacity
         v1group: InMemCollector
@@ -895,7 +898,24 @@ groups:
           1000x) of the total number of concurrently active traces (trace throughput *
           trace duration).
 
-      - name: MaxMemory
+      - name: AvailableMemory
+        type: memorysize
+        valuetype: memorysize
+        example: "4Gb"
+        reload: true
+        envvar: REFINERY_AVAILABLE_MEMORY
+        commandLine: AvailableMemory
+        summary: is the amount of system memory available to the refinery process.
+        description: >
+          The amount of system memory available to the refinery process. This
+          value will typically be set through an environment variable controlled
+          by the container or deploy script. If this value is zero or not set,
+          MaxMemory cannot be used to calculate the maximum allocation and
+          MaxAlloc will be used instead. If set, this must be a memory size.
+          64-bit values are supported. Sizes with standard unit suffixes like
+          "MB" and "GiB" are also supported.
+
+      - name: MaxMemoryPercentage
         type: percentage
         valuetype: showexample
         default: 75
@@ -905,20 +925,17 @@ groups:
           - type: minimum
             arg: 10
           - type: requiredWith
-            arg: MaxAlloc
+            arg: AvailableMemory
         summary: is the maximum percentage of memory that should be allocated by the span collector.
         description: >
-          If nonzero, it must be an integer value between 1 and 100, representing
-          the target maximum percentage of memory that should be allocated by
-          the span collector. If set to a non-zero value, once per tick (see
-          SendTicker) the collector will compare total allocated bytes to this
-          calculated value. If allocation is too high, traces will be ejected
-          from the cache early to reduce memory. Useful values for this setting
-          are generally in the range of 70-90. Depending on deployment details,
-          system memory information may not be available. If it is not, a
-          warning will be logged and the value of MaxAlloc will be used.
-          If this value is 0, MaxAlloc will be used. Requires MaxAlloc to be
-          nonzero.
+          If nonzero, it must be an integer value between 1 and 100,
+          representing the target maximum percentage of memory that should be
+          allocated by the span collector. If set to a non-zero value, once per
+          tick (see SendTicker) the collector will compare total allocated bytes
+          to this calculated value. If allocation is too high, traces will be
+          ejected from the cache early to reduce memory. Useful values for this
+          setting are generally in the range of 70-90. If this value is 0,
+          MaxAlloc will be used.
 
       - name: MaxAlloc
         v1group: InMemCollector
@@ -927,9 +944,6 @@ groups:
         valuetype: memorysize
         default: 0
         reload: true
-        validations:
-          - type: requiredWith
-            arg: MaxMemory
         summary: is the maximum number of bytes that should be allocated by the collector.
         description: >
           If set, this must be a memory size. 64-bit values are supported. Sizes
diff --git a/config/mock.go b/config/mock.go
index 9abb838..d6fc29d 100644
--- a/config/mock.go
+++ b/config/mock.go
@@ -9,78 +9,84 @@ import (
 // MockConfig will respond with whatever config it's set to do during
 // initialization
 type MockConfig struct {
-	Callbacks                            []func()
-	IsAPIKeyValidFunc                    func(string) bool
-	GetInMemoryCollectorCacheCapacityErr error
-	GetInMemoryCollectorCacheCapacityVal CollectionConfig
-	GetHoneycombAPIErr                   error
-	GetHoneycombAPIVal                   string
-	GetListenAddrErr                     error
-	GetListenAddrVal                     string
-	GetPeerListenAddrErr                 error
-	GetPeerListenAddrVal                 string
-	GetCompressPeerCommunicationsVal     bool
-	GetGRPCListenAddrErr                 error
-	GetGRPCListenAddrVal                 string
-	GetLoggerTypeErr                     error
-	GetLoggerTypeVal                     string
-	GetHoneycombLoggerConfigErr          error
-	GetHoneycombLoggerConfigVal          HoneycombLoggerConfig
-	GetLoggerLevelVal                    Level
-	GetPeersErr                          error
-	GetPeersVal                          []string
-	GetRedisHostErr                      error
-	GetRedisHostVal                      string
-	GetRedisUsernameErr                  error
-	GetRedisUsernameVal                  string
-	GetRedisPasswordErr                  error
-	GetRedisPasswordVal                  string
-	GetRedisDatabaseVal                  int
-	GetRedisPrefixVal                    string
-	GetUseTLSErr                         error
-	GetUseTLSVal                         bool
-	GetUseTLSInsecureErr                 error
-	GetUseTLSInsecureVal                 bool
-	GetSamplerTypeErr                    error
-	GetSamplerTypeName                   string
-	GetSamplerTypeVal                    interface{}
-	GetLegacyMetricsConfigVal            LegacyMetricsConfig
-	GetPrometheusMetricsConfigVal        PrometheusMetricsConfig
-	GetOTelMetricsConfigVal              OTelMetricsConfig
-	GetSendDelayErr                      error
-	GetSendDelayVal                      time.Duration
-	GetBatchTimeoutVal                   time.Duration
-	GetTraceTimeoutErr                   error
-	GetTraceTimeoutVal                   time.Duration
-	GetMaxBatchSizeVal                   uint
-	GetUpstreamBufferSizeVal             int
-	GetPeerBufferSizeVal                 int
-	SendTickerVal                        time.Duration
-	IdentifierInterfaceName              string
-	UseIPV6Identifier                    bool
-	RedisIdentifier                      string
-	PeerManagementType                   string
-	DebugServiceAddr                     string
-	DryRun                               bool
-	AddHostMetadataToTrace               bool
-	AddRuleReasonToTrace                 bool
-	EnvironmentCacheTTL                  time.Duration
-	DatasetPrefix                        string
-	QueryAuthToken                       string
-	GRPCMaxConnectionIdle                time.Duration
-	GRPCMaxConnectionAge                 time.Duration
-	GRPCMaxConnectionAgeGrace            time.Duration
-	GRPCTime                             time.Duration
-	GRPCTimeout                          time.Duration
-	PeerTimeout                          time.Duration
-	AdditionalErrorFields                []string
-	AddSpanCountToRoot                   bool
-	SampleCache                          SampleCacheConfig
-	StressRelief                         StressReliefConfig
-	AdditionalAttributes                 map[string]string
-	TraceIdFieldNames                    []string
-	ParentIdFieldNames                   []string
-	CfgMetadata                          []ConfigMetadata
+	Callbacks                        []func()
+	IsAPIKeyValidFunc                func(string) bool
+	GetCollectorTypeErr              error
+	GetCollectorTypeVal              string
+	GetCollectionConfigErr           error
+	GetCollectionConfigVal           CollectionConfig
+	GetHoneycombAPIErr               error
+	GetHoneycombAPIVal               string
+	GetListenAddrErr                 error
+	GetListenAddrVal                 string
+	GetPeerListenAddrErr             error
+	GetPeerListenAddrVal             string
+	GetCompressPeerCommunicationsVal bool
+	GetGRPCListenAddrErr             error
+	GetGRPCListenAddrVal             string
+	GetLoggerTypeErr                 error
+	GetLoggerTypeVal                 string
+	GetHoneycombLoggerConfigErr      error
+	GetHoneycombLoggerConfigVal      HoneycombLoggerConfig
+	GetLoggerLevelVal                Level
+	GetPeersErr                      error
+	GetPeersVal                      []string
+	GetRedisHostErr                  error
+	GetRedisHostVal                  string
+	GetRedisUsernameErr              error
+	GetRedisUsernameVal              string
+	GetRedisPasswordErr              error
+	GetRedisPasswordVal              string
+	GetRedisDatabaseVal              int
+	GetRedisPrefixVal                string
+	GetUseTLSErr                     error
+	GetUseTLSVal                     bool
+	GetUseTLSInsecureErr             error
+	GetUseTLSInsecureVal             bool
+	GetSamplerTypeErr                error
+	GetSamplerTypeName               string
+	GetSamplerTypeVal                interface{}
+	GetMetricsTypeErr                error
+	GetMetricsTypeVal                string
+	GetLegacyMetricsConfigVal        LegacyMetricsConfig
+	GetPrometheusMetricsConfigVal    PrometheusMetricsConfig
+	GetOTelMetricsConfigVal          OTelMetricsConfig
+	GetSendDelayErr                  error
+	GetSendDelayVal                  time.Duration
+	GetBatchTimeoutVal               time.Duration
+	GetTraceTimeoutErr               error
+	GetTraceTimeoutVal               time.Duration
+	GetMaxBatchSizeVal               uint
+	GetUpstreamBufferSizeVal         int
+	GetPeerBufferSizeVal             int
+	SendTickerVal                    time.Duration
+	IdentifierInterfaceName          string
+	UseIPV6Identifier                bool
+	RedisIdentifier                  string
+	PeerManagementType               string
+	DebugServiceAddr                 string
+	DryRun                           bool
+	DryRunFieldName                  string
+	AddHostMetadataToTrace           bool
+	AddRuleReasonToTrace             bool
+	EnvironmentCacheTTL              time.Duration
+	DatasetPrefix                    string
+	QueryAuthToken                   string
+	GRPCMaxConnectionIdle            time.Duration
+	GRPCMaxConnectionAge             time.Duration
+	GRPCMaxConnectionAgeGrace        time.Duration
+	GRPCTime                         time.Duration
+	GRPCTimeout                      time.Duration
+	PeerTimeout                      time.Duration
+	AdditionalErrorFields            []string
+	AddSpanCountToRoot               bool
+	CacheOverrunStrategy             string
+	SampleCache                      SampleCacheConfig
+	StressRelief                     StressReliefConfig
+	AdditionalAttributes             map[string]string
+	TraceIdFieldNames                []string
+	ParentIdFieldNames               []string
+	CfgMetadata                      []ConfigMetadata
 
 	Mux sync.RWMutex
 }
@@ -112,11 +118,18 @@ func (m *MockConfig) IsAPIKeyValid(key string) bool {
 	return m.IsAPIKeyValidFunc(key)
 }
 
-func (m *MockConfig) GetInMemCollectorCacheCapacity() (CollectionConfig, error) {
+func (m *MockConfig) GetCollectorType() (string, error) {
 	m.Mux.RLock()
 	defer m.Mux.RUnlock()
 
-	return m.GetInMemoryCollectorCacheCapacityVal, m.GetInMemoryCollectorCacheCapacityErr
+	return m.GetCollectorTypeVal, m.GetCollectorTypeErr
+}
+
+func (m *MockConfig) GetCollectionConfig() (CollectionConfig, error) {
+	m.Mux.RLock()
+	defer m.Mux.RUnlock()
+
+	return m.GetCollectionConfigVal, m.GetCollectionConfigErr
 }
 
 func (m *MockConfig) GetHoneycombAPI() (string, error) {
diff --git a/tools/convert/templates/configV2.tmpl b/tools/convert/templates/configV2.tmpl
index 1feade5..a07724b 100644
--- a/tools/convert/templates/configV2.tmpl
+++ b/tools/convert/templates/configV2.tmpl
@@ -2,7 +2,7 @@
 ## Honeycomb Refinery Configuration ##
 ######################################
 #
-# created {{ now }} from {{ .Input }} using a template generated on 2023-06-20 at 16:34:39 UTC
+# created {{ now }} from {{ .Input }} using a template generated on 2023-06-20 at 18:40:46 UTC
 
 # This file contains the configuration for the Honeycomb Refinery.
 # More stuff will go here later.
@@ -710,7 +710,10 @@ RedisPeerManagement:
 #########################
 Collection:
   ## Brings together the settings that are relevant to collecting spans
-  ## together to make traces.
+  ## together to make traces. If none of the memory settings are used, then
+  ## Refinery will not attempt to limit its memory usage. This is not
+  ## recommended for production use since a burst of traffic could cause
+  ## Refinery to run out of memory and crash.
   ##
   ####
     ## CacheCapacity is the number of traces to keep in the cache's circular
@@ -727,8 +730,22 @@ Collection:
     ## Eligible for live reload.
     {{ nonDefaultOnly .Data "CacheCapacity" "InMemCollector.CacheCapacity" 10000 }}
 
-    ## MaxMemory is the maximum percentage of memory that should be allocated
-    ## by the span collector.
+    ## AvailableMemory is the amount of system memory available to the
+    ## refinery process.
+    ##
+    ## The amount of system memory available to the refinery process. This
+    ## value will typically be set through an environment variable controlled
+    ## by the container or deploy script. If this value is zero or not set,
+    ## MaxMemory cannot be used to calculate the maximum allocation and
+    ## MaxAlloc will be used instead. If set, this must be a memory size.
+    ## 64-bit values are supported. Sizes with standard unit suffixes like
+    ## "MB" and "GiB" are also supported.
+    ##
+    ## Eligible for live reload.
+    {{ memorysize .Data "AvailableMemory" "AvailableMemory" "4Gb" }}
+
+    ## MaxMemoryPercentage is the maximum percentage of memory that should be
+    ## allocated by the span collector.
     ##
     ## If nonzero, it must be an integer value between 1 and 100,
     ## representing the target maximum percentage of memory that should be
@@ -736,15 +753,12 @@ Collection:
     ## tick (see SendTicker) the collector will compare total allocated bytes
     ## to this calculated value. If allocation is too high, traces will be
     ## ejected from the cache early to reduce memory. Useful values for this
-    ## setting are generally in the range of 70-90. Depending on deployment
-    ## details, system memory information may not be available. If it is not,
-    ## a warning will be logged and the value of MaxAlloc will be used. If
-    ## this value is 0, MaxAlloc will be used. Requires MaxAlloc to be
-    ## nonzero.
+    ## setting are generally in the range of 70-90. If this value is 0,
+    ## MaxAlloc will be used.
     ##
     ## default: 75
     ## Eligible for live reload.
-    # MaxMemory: 75
+    # MaxMemoryPercentage: 75
 
     ## MaxAlloc is the maximum number of bytes that should be allocated by
     ## the collector.


Test patch:
diff --git a/app/app_test.go b/app/app_test.go
index d865740..c3af3f0 100644
--- a/app/app_test.go
+++ b/app/app_test.go
@@ -103,23 +103,23 @@ func newStartedApp(
 	enableHostMetadata bool,
 ) (*App, inject.Graph) {
 	c := &config.MockConfig{
-		GetSendDelayVal:                      0,
-		GetTraceTimeoutVal:                   10 * time.Millisecond,
-		GetMaxBatchSizeVal:                   500,
-		GetSamplerTypeVal:                    &config.DeterministicSamplerConfig{SampleRate: 1},
-		SendTickerVal:                        2 * time.Millisecond,
-		PeerManagementType:                   "file",
-		GetUpstreamBufferSizeVal:             10000,
-		GetPeerBufferSizeVal:                 10000,
-		GetListenAddrVal:                     "127.0.0.1:" + strconv.Itoa(basePort),
-		GetPeerListenAddrVal:                 "127.0.0.1:" + strconv.Itoa(basePort+1),
-		IsAPIKeyValidFunc:                    func(k string) bool { return k == legacyAPIKey || k == nonLegacyAPIKey },
-		GetHoneycombAPIVal:                   "http://api.honeycomb.io",
-		GetInMemoryCollectorCacheCapacityVal: config.CollectionConfig{CacheCapacity: 10000},
-		AddHostMetadataToTrace:               enableHostMetadata,
-		TraceIdFieldNames:                    []string{"trace.trace_id"},
-		ParentIdFieldNames:                   []string{"trace.parent_id"},
-		SampleCache:                          config.SampleCacheConfig{KeptSize: 10000, DroppedSize: 100000, SizeCheckInterval: config.Duration(10 * time.Second)},
+		GetSendDelayVal:          0,
+		GetTraceTimeoutVal:       10 * time.Millisecond,
+		GetMaxBatchSizeVal:       500,
+		GetSamplerTypeVal:        &config.DeterministicSamplerConfig{SampleRate: 1},
+		SendTickerVal:            2 * time.Millisecond,
+		PeerManagementType:       "file",
+		GetUpstreamBufferSizeVal: 10000,
+		GetPeerBufferSizeVal:     10000,
+		GetListenAddrVal:         "127.0.0.1:" + strconv.Itoa(basePort),
+		GetPeerListenAddrVal:     "127.0.0.1:" + strconv.Itoa(basePort+1),
+		IsAPIKeyValidFunc:        func(k string) bool { return k == legacyAPIKey || k == nonLegacyAPIKey },
+		GetHoneycombAPIVal:       "http://api.honeycomb.io",
+		GetCollectionConfigVal:   config.CollectionConfig{CacheCapacity: 10000},
+		AddHostMetadataToTrace:   enableHostMetadata,
+		TraceIdFieldNames:        []string{"trace.trace_id"},
+		ParentIdFieldNames:       []string{"trace.parent_id"},
+		SampleCache:              config.SampleCacheConfig{KeptSize: 10000, DroppedSize: 100000, SizeCheckInterval: config.Duration(10 * time.Second)},
 	}
 
 	var err error
diff --git a/collect/collect_test.go b/collect/collect_test.go
index a426eeb..0ecca43 100644
--- a/collect/collect_test.go
+++ b/collect/collect_test.go
@@ -467,7 +467,7 @@ func TestCacheSizeReload(t *testing.T) {
 		GetTraceTimeoutVal: 10 * time.Minute,
 		GetSamplerTypeVal:  &config.DeterministicSamplerConfig{SampleRate: 1},
 		SendTickerVal:      2 * time.Millisecond,
-		GetInMemoryCollectorCacheCapacityVal: config.CollectionConfig{
+		GetCollectionConfigVal: config.CollectionConfig{
 			CacheCapacity: 1,
 		},
 		ParentIdFieldNames: []string{"trace.parent_id", "parentId"},
@@ -518,7 +518,7 @@ func TestCacheSizeReload(t *testing.T) {
 	assert.Eventually(t, check, 60*wait, wait, "expected one trace evicted and sent")
 
 	conf.Mux.Lock()
-	conf.GetInMemoryCollectorCacheCapacityVal.CacheCapacity = 2
+	conf.GetCollectionConfigVal.CacheCapacity = 2
 	conf.Mux.Unlock()
 	conf.ReloadConfig()
 
@@ -535,7 +535,7 @@ func TestCacheSizeReload(t *testing.T) {
 	assert.True(t, check(), "expected no more traces evicted and sent")
 
 	conf.Mux.Lock()
-	conf.GetInMemoryCollectorCacheCapacityVal.CacheCapacity = 1
+	conf.GetCollectionConfigVal.CacheCapacity = 1
 	conf.Mux.Unlock()
 	conf.ReloadConfig()
 
@@ -549,12 +549,12 @@ func TestSampleConfigReload(t *testing.T) {
 	transmission.Start()
 
 	conf := &config.MockConfig{
-		GetSendDelayVal:                      0,
-		GetTraceTimeoutVal:                   60 * time.Second,
-		GetSamplerTypeVal:                    &config.DeterministicSamplerConfig{SampleRate: 1},
-		SendTickerVal:                        2 * time.Millisecond,
-		ParentIdFieldNames:                   []string{"trace.parent_id", "parentId"},
-		GetInMemoryCollectorCacheCapacityVal: config.CollectionConfig{CacheCapacity: 10},
+		GetSendDelayVal:        0,
+		GetTraceTimeoutVal:     60 * time.Second,
+		GetSamplerTypeVal:      &config.DeterministicSamplerConfig{SampleRate: 1},
+		SendTickerVal:          2 * time.Millisecond,
+		ParentIdFieldNames:     []string{"trace.parent_id", "parentId"},
+		GetCollectionConfigVal: config.CollectionConfig{CacheCapacity: 10},
 		SampleCache: config.SampleCacheConfig{
 			KeptSize:          100,
 			DroppedSize:       100,
@@ -696,7 +696,7 @@ func TestStableMaxAlloc(t *testing.T) {
 	var mem runtime.MemStats
 	runtime.ReadMemStats(&mem)
 	// Set MaxAlloc, which should cause cache evictions.
-	conf.GetInMemoryCollectorCacheCapacityVal.MaxAlloc = config.MemorySize(mem.Alloc * 99 / 100)
+	conf.GetCollectionConfigVal.MaxAlloc = config.MemorySize(mem.Alloc * 99 / 100)
 
 	coll.mutex.Unlock()
 
diff --git a/config/config_test.go b/config/config_test.go
index 54aa379..270da03 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -405,7 +405,7 @@ func TestMaxAlloc(t *testing.T) {
 	assert.NoError(t, err)
 
 	expected := MemorySize(16 * 1024 * 1024 * 1024)
-	inMemConfig, err := c.GetInMemCollectorCacheCapacity()
+	inMemConfig, err := c.GetCollectionConfig()
 	assert.NoError(t, err)
 	assert.Equal(t, expected, inMemConfig.MaxAlloc)
 }
diff --git a/route/otlp_trace_test.go b/route/otlp_trace_test.go
index 4e19520..efdc672 100644
--- a/route/otlp_trace_test.go
+++ b/route/otlp_trace_test.go
@@ -59,7 +59,7 @@ func TestOTLPHandler(t *testing.T) {
 		GetTraceTimeoutVal: 60 * time.Second,
 		GetSamplerTypeVal:  &config.DeterministicSamplerConfig{SampleRate: 1},
 		SendTickerVal:      2 * time.Millisecond,
-		GetInMemoryCollectorCacheCapacityVal: config.CollectionConfig{
+		GetCollectionConfigVal: config.CollectionConfig{
 			CacheCapacity: 100,
 			MaxAlloc:      100,
 		},