From d6795b491c748855e4a4d6eec07ae2cb44a78da1 Mon Sep 17 00:00:00 2001
From: Mike Dame <mikedame@google.com>
Date: Wed, 18 Sep 2024 09:13:46 -0400
Subject: [PATCH] Support timeout in Collector Logs exporter (#894)

* Support timeout in Collector Logs exporter

* Add default timeout

* Add duration to test
---
 .../integrationtest/inmemoryotelexporter.go        |  2 ++
 .../integrationtest/logs_integration_test.go       |  3 +++
 exporter/collector/logs.go                         | 14 ++++++++++++--
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/exporter/collector/integrationtest/inmemoryotelexporter.go b/exporter/collector/integrationtest/inmemoryotelexporter.go
index c4c91016..1fc2c855 100644
--- a/exporter/collector/integrationtest/inmemoryotelexporter.go
+++ b/exporter/collector/integrationtest/inmemoryotelexporter.go
@@ -166,12 +166,14 @@ func NewLogTestExporter(
 	logger, err := zap.NewDevelopment()
 	require.NoError(t, err)
 
+	var duration time.Duration
 	exporter, err := collector.NewGoogleCloudLogsExporter(
 		ctx,
 		cfg,
 		logger,
 		meterProvider,
 		"latest",
+		duration,
 	)
 	require.NoError(t, err)
 
diff --git a/exporter/collector/integrationtest/logs_integration_test.go b/exporter/collector/integrationtest/logs_integration_test.go
index b1759d68..3dfbc173 100644
--- a/exporter/collector/integrationtest/logs_integration_test.go
+++ b/exporter/collector/integrationtest/logs_integration_test.go
@@ -43,12 +43,15 @@ func createLogsExporter(
 	cfg := test.CreateLogConfig()
 	// For sending to a real project, set the project ID from an env var.
 	cfg.ProjectID = os.Getenv("PROJECT_ID")
+
+	var duration time.Duration
 	exporter, err := collector.NewGoogleCloudLogsExporter(
 		ctx,
 		cfg,
 		logger,
 		noop.NewMeterProvider(),
 		"latest",
+		duration,
 	)
 	require.NoError(t, err)
 	err = exporter.Start(ctx, componenttest.NewNopHost())
diff --git a/exporter/collector/logs.go b/exporter/collector/logs.go
index b836078a..79763b4d 100644
--- a/exporter/collector/logs.go
+++ b/exporter/collector/logs.go
@@ -148,6 +148,7 @@ type LogsExporter struct {
 	loggingClient *loggingv2.Client
 	cfg           Config
 	mapper        logMapper
+	timeout       time.Duration
 }
 
 type logMapper struct {
@@ -163,6 +164,7 @@ func NewGoogleCloudLogsExporter(
 	log *zap.Logger,
 	meterProvider metric.MeterProvider,
 	version string,
+	timeout time.Duration,
 ) (*LogsExporter, error) {
 	setVersionInUserAgent(&cfg, version)
 	obs := selfObservability{
@@ -171,8 +173,9 @@ func NewGoogleCloudLogsExporter(
 	}
 
 	return &LogsExporter{
-		cfg: cfg,
-		obs: obs,
+		cfg:     cfg,
+		obs:     obs,
+		timeout: timeout,
 		mapper: logMapper{
 			obs:            obs,
 			cfg:            cfg,
@@ -364,6 +367,13 @@ func mergeLogLabels(instrumentationSource, instrumentationVersion string, resour
 }
 
 func (l *LogsExporter) writeLogEntries(ctx context.Context, batch []*logpb.LogEntry) (*logpb.WriteLogEntriesResponse, error) {
+	timeout := l.timeout
+	if timeout <= 0 {
+		timeout = 12 * time.Second
+	}
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
 	request := &logpb.WriteLogEntriesRequest{
 		PartialSuccess: true,
 		Entries:        batch,