cleanup from pr comments

ecs-agent/introspection/server.go

@@ -93,20 +93,23 @@ func NewServer(agentState v1.AgentState, metricsFactory metrics.EntryFactory, op
 	}
 	return setup(agentState, metricsFactory, config)
 }
+
 func v1HandlersSetup(serverMux *http.ServeMux,
 	agentState v1.AgentState,
 	metricsFactory metrics.EntryFactory) {
 	serverMux.HandleFunc(handlers.V1AgentMetadataPath, handlers.AgentMetadataHandler(agentState, metricsFactory))
 	serverMux.HandleFunc(handlers.V1TasksMetadataPath, handlers.TasksMetadataHandler(agentState, metricsFactory))
 	serverMux.HandleFunc(licensePath, licenseHandler(agentState, metricsFactory))
 }
+
 func pprofHandlerSetup(serverMux *http.ServeMux) {
 	serverMux.HandleFunc(pprofBasePath, pprofIndexHandler)
 	serverMux.HandleFunc(pprofCMDLinePath, pprofCmdlineHandler)
 	serverMux.HandleFunc(pprofProfilePath, pprofProfileHandler)
 	serverMux.HandleFunc(pprofSymbolPath, pprofSymbolHandler)
 	serverMux.HandleFunc(pprofTracePath, pprofTraceHandler)
 }
+
 func setup(
 	agentState v1.AgentState,
 	metricsFactory metrics.EntryFactory,

ecs-agent/introspection/server_integ_test.go

@@ -125,7 +125,7 @@ func TestWriteTimeout(t *testing.T) {
 				return testAgentMetadata, nil
 			})
 
-		response, err := performRequest(t, mockAgentState, mockMetricsFactory, handlers.V1AgentMetadataPath, WithReadTimeout(time.Millisecond*150))
+		response, err := performRequest(t, mockAgentState, mockMetricsFactory, handlers.V1AgentMetadataPath, WithWriteTimeout(time.Millisecond*150))
 		require.NoError(t, err)
 
 		bodyBytes, err := io.ReadAll(response.Body)

ecs-agent/introspection/server_test.go

@@ -30,6 +30,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+const (
+	writeTimeout = time.Second * 5
+	readTimeout  = time.Second * 6
+)
+
 var runtimeStatsConfigForTest = false
 
 func TestNewServerErrors(t *testing.T) {
@@ -46,6 +51,35 @@ func TestNewServerErrors(t *testing.T) {
 	})
 }
 
+func TestNewServerConfig(t *testing.T) {
+	t.Run("read/write timeouts", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		agentState := mock_v1.NewMockAgentState(ctrl)
+		metricsFactory := mock_metrics.NewMockEntryFactory(ctrl)
+		server, err := NewServer(agentState, metricsFactory,
+			WithReadTimeout(readTimeout),
+			WithWriteTimeout(writeTimeout),
+		)
+		assert.Nil(t, err)
+		assert.Equal(t, readTimeout, server.ReadTimeout)
+		assert.Equal(t, writeTimeout, server.WriteTimeout)
+	})
+
+	t.Run("overwrite write timeout when profiling enabled", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		agentState := mock_v1.NewMockAgentState(ctrl)
+		metricsFactory := mock_metrics.NewMockEntryFactory(ctrl)
+		server, err := NewServer(agentState, metricsFactory,
+			WithReadTimeout(readTimeout),
+			WithWriteTimeout(writeTimeout),
+			WithRuntimeStats(true),
+		)
+		assert.Nil(t, err)
+		assert.Equal(t, readTimeout, server.ReadTimeout)
+		assert.Equal(t, writeTimeoutForPprof, server.WriteTimeout)
+	})
+}
+
 func TestNewServerDefaults(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	agentState := mock_v1.NewMockAgentState(ctrl)

ecs-agent/introspection/v1/handlers/handlers.go

@@ -101,15 +101,15 @@ func TasksMetadataHandler(
 		if dockerIDExists {
 			// Find the task that has a container with a matching docker ID.
 			if len(dockerID) > dockerShortIDLen {
-				getTaskByID(agentState, metricsFactory, dockerID, w)
+				getTaskMetadata(agentState.GetTaskMetadataByID, dockerID, field.DockerId, metricsFactory, w)
 				return
 			} else {
-				getTaskByShortID(agentState, metricsFactory, dockerID, w)
+				getTaskMetadata(agentState.GetTaskMetadataByShortID, dockerID, field.DockerId, metricsFactory, w)
 				return
 			}
 		} else if taskARNExists {
 			// Find the task with a matching Arn.
-			getTaskByARN(agentState, metricsFactory, taskArn, w)
+			getTaskMetadata(agentState.GetTaskMetadataByArn, taskArn, field.TaskARN, metricsFactory, w)
 			return
 		} else {
 			// Return all tasks.
@@ -137,63 +137,18 @@ func getTasksMetadata(
 	tmdsutils.WriteJSONResponse(w, http.StatusOK, tasksMetadata, requestTypeTasks)
 }
 
-// getTaskByARN writes metadata for the corresponding task to the response, or an
-// error status if the agent cannot return the metadata
-func getTaskByARN(
-	agentState v1.AgentState,
-	metricsFactory metrics.EntryFactory,
-	taskARN string,
-	w http.ResponseWriter,
-) {
-	taskMetadata, err := agentState.GetTaskMetadataByArn(taskARN)
-	if err != nil {
-		logger.Error("Failed to get v1 task metadata.", logger.Fields{
-			field.Error:   err,
-			field.TaskARN: taskARN,
-		})
-		responseCode, metricName := getHTTPErrorCode(err)
-		metricsFactory.New(metricName).Done(err)
-		tmdsutils.WriteJSONResponse(w, responseCode, v1.TaskResponse{}, requestTypeTasks)
-		return
-	}
-	tmdsutils.WriteJSONResponse(w, http.StatusOK, taskMetadata, requestTypeTasks)
-}
-
-// getTaskByID writes metadata for the corresponding task to the response, or an
-// error status if the agent cannot return the metadata
-func getTaskByID(
-	agentState v1.AgentState,
+func getTaskMetadata(
+	lookupFn func(key string) (*v1.TaskResponse, error),
+	key string,
+	keyType string,
 	metricsFactory metrics.EntryFactory,
-	dockerID string,
 	w http.ResponseWriter,
 ) {
-	taskMetadata, err := agentState.GetTaskMetadataByID(dockerID)
+	taskMetadata, err := lookupFn(key)
 	if err != nil {
 		logger.Error("Failed to get v1 task metadata.", logger.Fields{
-			field.Error:    err,
-			field.DockerId: dockerID,
-		})
-		responseCode, metricName := getHTTPErrorCode(err)
-		metricsFactory.New(metricName).Done(err)
-		tmdsutils.WriteJSONResponse(w, responseCode, v1.TaskResponse{}, requestTypeTasks)
-		return
-	}
-	tmdsutils.WriteJSONResponse(w, http.StatusOK, taskMetadata, requestTypeTasks)
-}
-
-// getTaskByShortID writes metadata for the corresponding task to the response, or an
-// error status if the agent cannot return the metadata
-func getTaskByShortID(
-	agentState v1.AgentState,
-	metricsFactory metrics.EntryFactory,
-	shortDockerID string,
-	w http.ResponseWriter,
-) {
-	taskMetadata, err := agentState.GetTaskMetadataByShortID(shortDockerID)
-	if err != nil {
-		logger.Error("Failed to get v1 task metadata.", logger.Fields{
-			field.Error:    err,
-			field.DockerId: shortDockerID,
+			field.Error: err,
+			keyType:     key,
 		})
 		responseCode, metricName := getHTTPErrorCode(err)
 		metricsFactory.New(metricName).Done(err)