fix: isolate exporter timeout phases

internal/exporter/collector/collector.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"crypto/rand"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"time"
 
@@ -124,33 +125,55 @@ func (c *collector) Collect(ctx context.Context) (*HealthData, error) {
 		Timestamp:    time.Now().UTC(),
 	}
 
+	if err := ctx.Err(); err != nil {
+		return data, err
+	}
+
 	// Collect machine info if enabled
 	if c.config.IncludeMachineInfo {
 		if err := c.collectMachineInfo(data); err != nil {
 			log.Logger.Errorw("Failed to collect machine info", "error", err)
 		}
 	}
+	if err := ctx.Err(); err != nil {
+		return data, err
+	}
 
 	// Collect metrics if enabled
 	if c.config.IncludeMetrics {
 		if err := c.collectMetrics(ctx, data); err != nil {
+			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+				return data, err
+			}
 			log.Logger.Errorw("Failed to collect metrics", "error", err)
 		}
 	}
+	if err := ctx.Err(); err != nil {
+		return data, err
+	}
 
 	// Collect events if enabled
 	if c.config.IncludeEvents {
 		if err := c.collectEvents(ctx, data); err != nil {
+			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+				return data, err
+			}
 			log.Logger.Errorw("Failed to collect events", "error", err)
 		}
 	}
+	if err := ctx.Err(); err != nil {
+		return data, err
+	}
 
 	// Collect component data if enabled
 	if c.config.IncludeComponentData {
 		if err := c.collectComponentData(data); err != nil {
 			log.Logger.Errorw("Failed to collect component data", "error", err)
 		}
 	}
+	if err := ctx.Err(); err != nil {
+		return data, err
+	}
 
 	// Collect attestation data if provider is available
 	// Attestation is always enabled if manager is available
@@ -219,8 +242,17 @@ func (c *collector) collectEvents(ctx context.Context, data *HealthData) error {
 	}
 
 	for _, component := range components {
+		if err := ctx.Err(); err != nil {
+			data.Events = allEvents
+			return err
+		}
+
 		componentEvents, err := component.Events(ctx, since)
 		if err != nil {
+			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+				data.Events = allEvents
+				return err
+			}
 			log.Logger.Errorw("Failed to get events from component",
 				"component", component.Name(), "error", err)
 			continue

internal/exporter/collector/collector_test.go

@@ -378,6 +378,60 @@ func TestNew(t *testing.T) {
 	var _ = c
 }
 
+func TestCollectReturnsPartialDataWhenEventCollectionTimesOut(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	firstEvent := apiv1.Event{
+		Component: "cpu",
+		Time:      metav1.NewTime(time.Now().UTC()),
+		Name:      "test_event",
+		Type:      apiv1.EventTypeWarning,
+		Message:   "first event",
+	}
+
+	registry := &mockRegistry{
+		components: []components.Component{
+			&mockComponent{
+				name:   "cpu",
+				events: []apiv1.Event{firstEvent},
+			},
+			&mockComponent{
+				name: "memory",
+				eventFunc: func(ctx context.Context, since time.Time) (apiv1.Events, error) {
+					cancel()
+					<-ctx.Done()
+					return nil, ctx.Err()
+				},
+			},
+		},
+	}
+
+	c := New(
+		&config.HealthExporterConfig{
+			IncludeEvents: true,
+			EventsLookback: metav1.Duration{
+				Duration: time.Minute,
+			},
+		},
+		nil,
+		nil,
+		nil,
+		&mockEventStore{},
+		registry,
+		nil,
+		nil,
+		"test-machine-id",
+		nil,
+	)
+
+	data, err := c.Collect(ctx)
+	require.ErrorIs(t, err, context.Canceled)
+	require.NotNil(t, data)
+	require.Len(t, data.Events, 1)
+	assert.Equal(t, "cpu", data.Events[0].Component)
+	assert.Equal(t, "test_event", data.Events[0].Name)
+}
+
 func TestCollector_Collect_BasicFlow(t *testing.T) {
 	ctx := context.Background()
 	cfg := &config.HealthExporterConfig{
@@ -843,13 +897,17 @@ type mockComponent struct {
 	events       []apiv1.Event
 	healthStates []apiv1.HealthState
 	shouldError  bool
+	eventFunc    func(ctx context.Context, since time.Time) (apiv1.Events, error)
 }
 
 func (m *mockComponent) Name() string {
 	return m.name
 }
 
 func (m *mockComponent) Events(ctx context.Context, since time.Time) (apiv1.Events, error) {
+	if m.eventFunc != nil {
+		return m.eventFunc(ctx, since)
+	}
 	if m.shouldError {
 		return nil, errors.New("mock component error")
 	}

internal/exporter/exporter.go

@@ -24,6 +24,7 @@ package exporter
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"
 
@@ -173,26 +174,35 @@ func (e *healthExporter) ExportNow(ctx context.Context) error {
 // export performs the actual data export operation
 func (e *healthExporter) export() error {
 	log.Logger.Infow("Starting health export")
-	ctx, cancel := context.WithTimeout(e.ctx, e.options.timeout)
-	defer cancel()
+	collectCtx, cancelCollect := context.WithTimeout(e.ctx, e.options.timeout)
+	defer cancelCollect()
 
 	// Refresh configuration from metadata on every export
 	// If the endpoints/auth token are not empty, export will continue
 	// If the endpoints/auth token are empty, exportHTTP will skip
-	e.refreshConfigFromMetadata(ctx)
+	e.refreshConfigFromMetadata(collectCtx)
 
 	// Collect health data
-	healthData, err := e.collector.Collect(ctx)
-	if err != nil {
+	healthData, err := e.collector.Collect(collectCtx)
+	if err != nil && healthData == nil {
 		return fmt.Errorf("collection failed: %w", err)
 	}
+	if err != nil {
+		log.Logger.Warnw("Collection completed with partial data",
+			"error", err,
+			"timed_out", errors.Is(err, context.DeadlineExceeded),
+			"canceled", errors.Is(err, context.Canceled))
+	}
 
 	// Export data based on mode
 	if e.options.config.OfflineMode {
 		return e.exportToFile(healthData)
-	} else {
-		return e.exportToHTTP(ctx, healthData)
 	}
+
+	uploadCtx, cancelUpload := context.WithTimeout(e.ctx, e.options.timeout)
+	defer cancelUpload()
+
+	return e.exportToHTTP(uploadCtx, healthData)
 }
 
 // exportToFile writes health data to files

internal/exporter/exporter_test.go

@@ -40,6 +40,7 @@ import (
 
 	"github.com/NVIDIA/fleet-intelligence-agent/internal/config"
 	"github.com/NVIDIA/fleet-intelligence-agent/internal/exporter/collector"
+	"github.com/NVIDIA/fleet-intelligence-agent/internal/exporter/writer"
 	"github.com/NVIDIA/fleet-intelligence-agent/internal/machineinfo"
 )
 
@@ -58,6 +59,16 @@ func (m *MockCollector) Collect(ctx context.Context) (*collector.HealthData, err
 	return args.Get(0).(*collector.HealthData), args.Error(1)
 }
 
+type mockHTTPWriter struct {
+	sendFunc func(ctx context.Context, data *collector.HealthData, metricsEndpoint string, logsEndpoint string, maxRetries int, authToken string) (string, error)
+}
+
+func (m *mockHTTPWriter) Send(ctx context.Context, data *collector.HealthData, metricsEndpoint string, logsEndpoint string, maxRetries int, authToken string) (string, error) {
+	return m.sendFunc(ctx, data, metricsEndpoint, logsEndpoint, maxRetries, authToken)
+}
+
+func (m *mockHTTPWriter) SetJWTRefreshFunc(refreshFunc writer.JWTRefreshFunc) {}
+
 // MockMetricsStore is a mock implementation of pkgmetrics.Store
 type MockMetricsStore struct {
 	mock.Mock
@@ -1115,6 +1126,53 @@ func TestExportWithCollectorError(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestExportUsesFreshUploadContextAfterPartialCollection(t *testing.T) {
+	cfg := &config.HealthExporterConfig{
+		Interval:        metav1.Duration{Duration: 1 * time.Minute},
+		Timeout:         metav1.Duration{Duration: 20 * time.Millisecond},
+		MetricsEndpoint: "https://metrics.example.com",
+		LogsEndpoint:    "https://logs.example.com",
+		AuthToken:       "test-token",
+	}
+
+	ctx := context.Background()
+	exporter, err := New(ctx, WithConfig(cfg), WithMachineID("test-machine-id"))
+	require.NoError(t, err)
+
+	he := exporter.(*healthExporter)
+	healthData := &collector.HealthData{
+		CollectionID: "test-collection",
+		MachineID:    "test-machine-id",
+		Timestamp:    time.Now(),
+	}
+
+	mockCollector := &MockCollector{}
+	mockCollector.
+		On("Collect", mock.Anything).
+		Run(func(args mock.Arguments) {
+			<-args.Get(0).(context.Context).Done()
+		}).
+		Return(healthData, context.DeadlineExceeded)
+	he.collector = mockCollector
+
+	sendCalled := false
+	he.httpWriter = &mockHTTPWriter{
+		sendFunc: func(ctx context.Context, data *collector.HealthData, metricsEndpoint string, logsEndpoint string, maxRetries int, authToken string) (string, error) {
+			sendCalled = true
+			assert.NoError(t, ctx.Err(), "upload should use a fresh context")
+			assert.Equal(t, healthData, data)
+			return "", nil
+		},
+	}
+
+	err = he.export()
+	require.NoError(t, err)
+	assert.True(t, sendCalled, "upload should still run with partial collection data")
+
+	err = exporter.Stop()
+	require.NoError(t, err)
+}
+
 // TestExportHTTPWithReturnedToken tests handling of JWT token returned from server
 func TestExportHTTPWithReturnedToken(t *testing.T) {
 	newJWT := "new-jwt-token-from-server"