chore: add agentapi tests (#11269)

coderd/agentapi/metadata.go

@@ -12,6 +12,7 @@ import (
 	"cdr.dev/slog"
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 )
 
@@ -20,14 +21,26 @@ type MetadataAPI struct {
 	Database database.Store
 	Pubsub   pubsub.Pubsub
 	Log      slog.Logger
+
+	TimeNowFn func() time.Time // defaults to dbtime.Now()
+}
+
+func (a *MetadataAPI) now() time.Time {
+	if a.TimeNowFn != nil {
+		return a.TimeNowFn()
+	}
+	return dbtime.Now()
 }
 
 func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.BatchUpdateMetadataRequest) (*agentproto.BatchUpdateMetadataResponse, error) {
 	const (
-		// maxValueLen is set to 2048 to stay under the 8000 byte Postgres
-		// NOTIFY limit. Since both value and error can be set, the real payload
-		// limit is 2 * 2048 * 4/3 <base64 expansion> = 5461 bytes + a few
-		// hundred bytes for JSON syntax, key names, and metadata.
+		// maxAllKeysLen is the maximum length of all metadata keys. This is
+		// 6144 to stay below the Postgres NOTIFY limit of 8000 bytes, with some
+		// headway for the timestamp and JSON encoding. Any values that would
+		// exceed this limit are discarded (the rest are still inserted) and an
+		// error is returned.
+		maxAllKeysLen = 6144 // 1024 * 6
+
 		maxValueLen = 2048
 		maxErrorLen = maxValueLen
 	)
@@ -37,18 +50,36 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
 		return nil, err
 	}
 
-	collectedAt := time.Now()
-	dbUpdate := database.UpdateWorkspaceAgentMetadataParams{
-		WorkspaceAgentID: workspaceAgent.ID,
-		Key:              make([]string, 0, len(req.Metadata)),
-		Value:            make([]string, 0, len(req.Metadata)),
-		Error:            make([]string, 0, len(req.Metadata)),
-		CollectedAt:      make([]time.Time, 0, len(req.Metadata)),
-	}
-
+	var (
+		collectedAt = a.now()
+		allKeysLen  = 0
+		dbUpdate    = database.UpdateWorkspaceAgentMetadataParams{
+			WorkspaceAgentID: workspaceAgent.ID,
+			// These need to be `make(x, 0, len(req.Metadata))` instead of
+			// `make(x, len(req.Metadata))` because we may not insert all
+			// metadata if the keys are large.
+			Key:         make([]string, 0, len(req.Metadata)),
+			Value:       make([]string, 0, len(req.Metadata)),
+			Error:       make([]string, 0, len(req.Metadata)),
+			CollectedAt: make([]time.Time, 0, len(req.Metadata)),
+		}
+	)
 	for _, md := range req.Metadata {
 		metadataError := md.Result.Error
 
+		allKeysLen += len(md.Key)
+		if allKeysLen > maxAllKeysLen {
+			// We still insert the rest of the metadata, and we return an error
+			// after the insert.
+			a.Log.Warn(
+				ctx, "discarded extra agent metadata due to excessive key length",
+				slog.F("collected_at", collectedAt),
+				slog.F("all_keys_len", allKeysLen),
+				slog.F("max_all_keys_len", maxAllKeysLen),
+			)
+			break
+		}
+
 		// We overwrite the error if the provided payload is too long.
 		if len(md.Result.Value) > maxValueLen {
 			metadataError = fmt.Sprintf("value of %d bytes exceeded %d bytes", len(md.Result.Value), maxValueLen)
@@ -71,12 +102,16 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
 		a.Log.Debug(
 			ctx, "accepted metadata report",
 			slog.F("collected_at", collectedAt),
-			slog.F("original_collected_at", collectedAt),
 			slog.F("key", md.Key),
 			slog.F("value", ellipse(md.Result.Value, 16)),
 		)
 	}
 
+	err = a.Database.UpdateWorkspaceAgentMetadata(ctx, dbUpdate)
+	if err != nil {
+		return nil, xerrors.Errorf("update workspace agent metadata in database: %w", err)
+	}
+
 	payload, err := json.Marshal(WorkspaceAgentMetadataChannelPayload{
 		CollectedAt: collectedAt,
 		Keys:        dbUpdate.Key,
@@ -84,17 +119,17 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
 	if err != nil {
 		return nil, xerrors.Errorf("marshal workspace agent metadata channel payload: %w", err)
 	}
-
-	err = a.Database.UpdateWorkspaceAgentMetadata(ctx, dbUpdate)
-	if err != nil {
-		return nil, xerrors.Errorf("update workspace agent metadata in database: %w", err)
-	}
-
 	err = a.Pubsub.Publish(WatchWorkspaceAgentMetadataChannel(workspaceAgent.ID), payload)
 	if err != nil {
 		return nil, xerrors.Errorf("publish workspace agent metadata: %w", err)
 	}
 
+	// If the metadata keys were too large, we return an error so the agent can
+	// log it.
+	if allKeysLen > maxAllKeysLen {
+		return nil, xerrors.Errorf("metadata keys of %d bytes exceeded %d bytes", allKeysLen, maxAllKeysLen)
+	}
+
 	return &agentproto.BatchUpdateMetadataResponse{}, nil
 }