Merge pull request #75 from opensciencegrid/copilot/sub-pr-73

Address PR review feedback: config validation, test race safety, and documentation

Author: djw8605

README.md

@@ -189,8 +189,12 @@ See [config-collector.yaml](config/config-collector.yaml) for a complete example
 * `SHOVELER_INPUT_FOLLOW` - Follow file mode like tail: `true` or `false`
 
 **State Management (Collector Mode):**
-* `SHOVELER_STATE_ENTRY_TTL` - TTL for state entries in seconds (default: `300`)
-* `SHOVELER_STATE_MAX_ENTRIES` - Maximum state entries, 0 for unlimited (default: `0`)
+* `COLLECTOR_STATE_ENTRY_TTL` - TTL for state entries in seconds (default: `300`)
+* `COLLECTOR_STATE_MAX_ENTRIES` - Maximum state entries, 0 for unlimited (default: `0`)
+* `COLLECTOR_STATE_ENABLE_DNS_ENRICHMENT` - Enable DNS enrichment of monitoring records: `true` or `false` (default: `false`)
+* `COLLECTOR_STATE_DNS_CACHE_TTL` - DNS cache time-to-live in seconds (default: `3600`)
+* `COLLECTOR_STATE_DNS_WORKERS` - Number of concurrent DNS worker routines (default: `5`)
+* `COLLECTOR_STATE_DNS_TIMEOUT` - DNS query timeout in seconds (default: `2`)
 
 **Output Configuration (Collector Mode):**
 * `SHOVELER_OUTPUT_TYPE` - Output destination: `mq`, `file`, or `both` (default: `mq`)

collector/correlator.go

@@ -187,14 +187,14 @@ func NewCorrelatorWithConfig(config CorrelatorConfig) *Correlator {
 		config.Logger = logrus.New()
 	}
 
-	// Set DNS enrichment defaults
-	if config.DNSCacheTTL == 0 {
+	// Set DNS enrichment defaults (treat non-positive values as unset)
+	if config.DNSCacheTTL <= 0 {
 		config.DNSCacheTTL = 1 * time.Hour // Default 1 hour cache
 	}
-	if config.DNSWorkers == 0 {
+	if config.DNSWorkers <= 0 {
 		config.DNSWorkers = 5 // Default 5 concurrent DNS lookups
 	}
-	if config.DNSTimeout == 0 {
+	if config.DNSTimeout <= 0 {
 		config.DNSTimeout = 2 * time.Second // Default 2 second timeout
 	}
 

collector/dns_enrichment_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"strconv"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -15,11 +16,11 @@ import (
 // mockDNSResolver implements DNSResolver for testing
 type mockDNSResolver struct {
 	lookupFunc  func(ctx context.Context, addr string) ([]string, error)
-	lookupCount int
+	lookupCount atomic.Int64
 }
 
 func (m *mockDNSResolver) LookupAddr(ctx context.Context, addr string) ([]string, error) {
-	m.lookupCount++
+	m.lookupCount.Add(1)
 	if m.lookupFunc != nil {
 		return m.lookupFunc(ctx, addr)
 	}
@@ -84,7 +85,7 @@ func TestDNSEnrichment_CacheHit(t *testing.T) {
 	assert.False(t, needsAsync, "Should not need async on cache hit")
 
 	// Verify no DNS lookup was performed (cache hit)
-	assert.Equal(t, 0, mockResolver.lookupCount, "DNS lookup should not be called on cache hit")
+	assert.Equal(t, int64(0), mockResolver.lookupCount.Load(), "DNS lookup should not be called on cache hit")
 }
 
 // TestDNSEnrichment_CacheMiss_Success tests cache miss with successful lookup
@@ -122,19 +123,19 @@ func TestDNSEnrichment_CacheMiss_Success(t *testing.T) {
 	assert.Equal(t, "resolved.example.com", result, "Should return resolved hostname")
 
 	// Verify DNS lookup was performed
-	assert.Equal(t, 1, mockResolver.lookupCount, "DNS lookup should be called on cache miss")
+	assert.Equal(t, int64(1), mockResolver.lookupCount.Load(), "DNS lookup should be called on cache miss")
 
 	// Verify result was cached
 	val, exists := c.dnsCache.Get("192.0.2.2")
 	assert.True(t, exists, "Result should be cached")
 	assert.Equal(t, "resolved.example.com", val.(string), "Cached value should match resolved hostname")
 
 	// Second call should hit cache (no additional lookup)
-	mockResolver.lookupCount = 0
+	mockResolver.lookupCount.Store(0)
 	hostname2, needsAsync2 := c.enrichWithDNSSync("192.0.2.2")
 	assert.Equal(t, "resolved.example.com", hostname2, "Should return cached hostname")
 	assert.False(t, needsAsync2, "Should not need async on cache hit")
-	assert.Equal(t, 0, mockResolver.lookupCount, "Second call should hit cache")
+	assert.Equal(t, int64(0), mockResolver.lookupCount.Load(), "Second call should hit cache")
 }
 
 // TestDNSEnrichment_CacheMiss_Timeout tests cache miss with timeout
@@ -173,7 +174,7 @@ func TestDNSEnrichment_CacheMiss_Timeout(t *testing.T) {
 	assert.Equal(t, "", result, "Should return empty on timeout")
 
 	// Verify DNS lookup was attempted
-	assert.Equal(t, 1, mockResolver.lookupCount, "DNS lookup should be attempted")
+	assert.Equal(t, int64(1), mockResolver.lookupCount.Load(), "DNS lookup should be attempted")
 }
 
 // TestDNSEnrichment_CacheMiss_Failure tests cache miss with DNS failure
@@ -206,7 +207,7 @@ func TestDNSEnrichment_CacheMiss_Failure(t *testing.T) {
 	assert.Equal(t, "", result, "Should return empty on DNS failure")
 
 	// Verify DNS lookup was attempted
-	assert.Equal(t, 1, mockResolver.lookupCount, "DNS lookup should be attempted")
+	assert.Equal(t, int64(1), mockResolver.lookupCount.Load(), "DNS lookup should be attempted")
 }
 
 // TestDNSEnrichment_CacheTTL tests that cache entries expire after TTL
@@ -239,22 +240,27 @@ func TestDNSEnrichment_CacheTTL(t *testing.T) {
 	// First lookup - cache miss (blocking)
 	result1 := c.enrichWithDNSBlocking("192.0.2.5")
 	assert.Equal(t, "ttl-test.example.com", result1)
-	assert.Equal(t, 1, mockResolver.lookupCount)
+	assert.Equal(t, int64(1), mockResolver.lookupCount.Load())
 
 	// Immediate second lookup - cache hit
-	mockResolver.lookupCount = 0
+	mockResolver.lookupCount.Store(0)
 	hostname2, needsAsync2 := c.enrichWithDNSSync("192.0.2.5")
 	assert.Equal(t, "ttl-test.example.com", hostname2)
 	assert.False(t, needsAsync2, "Should hit cache")
-	assert.Equal(t, 0, mockResolver.lookupCount, "Should hit cache")
+	assert.Equal(t, int64(0), mockResolver.lookupCount.Load(), "Should hit cache")
 
-	// Wait for cache entry to expire
-	time.Sleep(600 * time.Millisecond)
+	// Wait for cache entry to expire and verify a new lookup occurs
+	assert.Eventually(t, func() bool {
+		mockResolver.lookupCount.Store(0)
+		_, needsAsync := c.enrichWithDNSSync("192.0.2.5")
+		return needsAsync
+	}, 5*time.Second, 100*time.Millisecond, "Cache entry should eventually expire")
 
-	// Third lookup after expiry - cache miss again
+	// Now do a blocking lookup to get fresh result
+	mockResolver.lookupCount.Store(0)
 	result3 := c.enrichWithDNSBlocking("192.0.2.5")
 	assert.Equal(t, "ttl-test.example.com", result3)
-	assert.Equal(t, 1, mockResolver.lookupCount, "Should do new lookup after TTL expiry")
+	assert.Equal(t, int64(1), mockResolver.lookupCount.Load(), "Should do new lookup after TTL expiry")
 }
 
 // TestDNSEnrichment_Concurrency tests multiple concurrent DNS lookups
@@ -306,7 +312,7 @@ func TestDNSEnrichment_Concurrency(t *testing.T) {
 
 	// All requests should succeed (workers handle them concurrently)
 	assert.Equal(t, numRequests, successCount, "All concurrent lookups should succeed")
-	assert.GreaterOrEqual(t, mockResolver.lookupCount, numRequests, "All lookups should be performed")
+	assert.GreaterOrEqual(t, mockResolver.lookupCount.Load(), int64(numRequests), "All lookups should be performed")
 }
 
 // TestDNSEnrichment_Integration tests full integration with record creation
@@ -426,7 +432,7 @@ func TestDNSEnrichment_Integration(t *testing.T) {
 	assert.Equal(t, "university.edu", record.UserDomain, "User domain should be extracted from enriched hostname")
 
 	// Verify DNS lookup was performed (once for user IP, once for server IP)
-	assert.Equal(t, 2, mockResolver.lookupCount, "DNS lookup should be performed for user and server IPs")
+	assert.Equal(t, int64(2), mockResolver.lookupCount.Load(), "DNS lookup should be performed for user and server IPs")
 
 	// Verify result was cached for user IP
 	val, exists := c.dnsCache.Get("192.0.2.10")

collector/integration_dns_test.go

@@ -3,7 +3,6 @@ package collector
 import (
 	"context"
 	"net"
-	"sync"
 	"testing"
 	"time"
 
@@ -136,26 +135,30 @@ func TestEndToEnd_DNSEnrichmentFlow(t *testing.T) {
 	// Verify the record needs DNS enrichment
 	assert.True(t, record.NeedsDNSEnrichment(), "Record should need DNS enrichment")
 
-	// Simulate what the main.go publish loop does
-	var wg sync.WaitGroup
-	wg.Add(1)
+	// Simulate what the main.go publish loop does with a timeout to avoid hanging
+	doneCh := make(chan struct{})
 
 	var enrichedRecord *CollectorRecord
 	correlator.EnrichRecordAsync(record, func(r *CollectorRecord) {
 		enrichedRecord = r
-		wg.Done()
+		close(doneCh)
 	})
 
-	// Wait for enrichment to complete
-	wg.Wait()
+	// Wait for enrichment to complete with timeout
+	select {
+	case <-doneCh:
+		// Success
+	case <-time.After(10 * time.Second):
+		t.Fatal("Timeout waiting for DNS enrichment to complete")
+	}
 
 	// Verify the enriched record has the user domain set
 	assert.NotNil(t, enrichedRecord)
 	assert.Equal(t, "example.com", enrichedRecord.UserDomain, "UserDomain should be extracted from DNS result")
 	assert.Equal(t, "server.example.com", enrichedRecord.ServerHostname, "ServerHostname should be resolved via DNS")
 
 	// Verify DNS was called for both user IP and server IP
-	assert.Equal(t, 2, mockResolver.lookupCount, "DNS resolver should have been called twice")
+	assert.Equal(t, int64(2), mockResolver.lookupCount.Load(), "DNS resolver should have been called twice")
 }
 
 // TestEndToEnd_DNSEnrichmentCacheHit tests that cached DNS results don't trigger async enrichment
@@ -280,5 +283,5 @@ func TestEndToEnd_DNSEnrichmentCacheHit(t *testing.T) {
 	assert.Equal(t, "server.example.com", record.ServerHostname, "ServerHostname should be set from cache")
 
 	// Verify DNS was NOT called (cache hit)
-	assert.Equal(t, 0, mockResolver.lookupCount, "DNS resolver should NOT have been called (cache hit)")
+	assert.Equal(t, int64(0), mockResolver.lookupCount.Load(), "DNS resolver should NOT have been called (cache hit)")
 }

config.go

@@ -145,10 +145,19 @@ func (c *Config) ReadConfigWithPathAndPrefix(configPath string, envPrefix string
 	c.State.EnableDNSEnrichment = viper.GetBool("state.enable_dns_enrichment")
 	viper.SetDefault("state.dns_cache_ttl", 3600) // 1 hour default
 	c.State.DNSCacheTTL = viper.GetInt("state.dns_cache_ttl")
+	if c.State.DNSCacheTTL <= 0 {
+		c.State.DNSCacheTTL = 3600
+	}
 	viper.SetDefault("state.dns_workers", 5) // 5 workers default
 	c.State.DNSWorkers = viper.GetInt("state.dns_workers")
+	if c.State.DNSWorkers <= 0 {
+		c.State.DNSWorkers = 5
+	}
 	viper.SetDefault("state.dns_timeout", 2) // 2 seconds default
 	c.State.DNSTimeout = viper.GetInt("state.dns_timeout")
+	if c.State.DNSTimeout <= 0 {
+		c.State.DNSTimeout = 2
+	}
 
 	// Output configuration (for collector mode)
 	viper.SetDefault("output.type", "mq") // message queue by default

config/config-collector.yaml

@@ -22,6 +22,11 @@ state:
   entry_ttl: 300
   # Maximum number of state entries (0 for unlimited)
   max_entries: 10000
+  # DNS enrichment settings (replaces reverse DNS lookup)
+  # enable_dns_enrichment: false
+  # dns_cache_ttl: 3600       # Time-to-live for DNS cache entries in seconds
+  # dns_workers: 5            # Number of worker goroutines for DNS lookups
+  # dns_timeout: 2            # DNS lookup timeout in seconds
 
 # Select which protocol to use in order to connect to the MQ
 # mq: amqp/stomp