buncha stuff

Author: michaelkedar

deployment/clouddeploy/gke-workers/base/importer.yaml

@@ -25,6 +25,10 @@ spec:
             env:
               - name: GITTER_HOST
                 value: http://gitter-service:8888
+              - name: IMPORT_TRACE_SAMPLE_RATE # for the overall importer trace
+                value: "1.0"
+              - name: TRACE_SAMPLE_RATE # for the individual vulnerability entries
+                value: "0.05"
             securityContext:
               privileged: true
             resources:

go/cmd/importer/main.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"strconv"
 	"syscall"
 	"time"
 
@@ -19,10 +20,18 @@ import (
 	"github.com/google/osv.dev/go/logger"
 	"github.com/google/osv.dev/go/osv/clients"
 	"github.com/hashicorp/go-retryablehttp"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
+	"google.golang.org/api/option"
 )
 
 func main() {
 	logger.InitGlobalLogger()
+	defer logger.Close()
+	ctx, span := otel.Tracer("importer").Start(context.Background(), "importer",
+		trace.WithAttributes(attribute.Float64("override_sample_rate", importerSampleRate())))
+	defer span.End()
 
 	strictValidation := flag.Bool("strict-validation", false, "Fail to import entries that do not pass validation. "+
 		"Note: this only applies to SourceRepositories with strict_validation=true")
@@ -35,14 +44,15 @@ func main() {
 
 	project := os.Getenv("GOOGLE_CLOUD_PROJECT")
 	if project == "" {
-		logger.Fatal("GOOGLE_CLOUD_PROJECT environment variable is not set")
+		logger.FatalContext(ctx, "GOOGLE_CLOUD_PROJECT environment variable is not set")
 	}
 
 	config := importer.Config{
 		StrictValidation: *strictValidation,
 		DeleteThreshold:  *deleteThresholdPct,
 		NumWorkers:       *numWorkers,
 		GitWorkDir:       filepath.Join(*workDir, "sources"),
+		SampleRate:       vulnerabilitySampleRate(),
 	}
 
 	httpClient := retryablehttp.NewClient()
@@ -52,36 +62,61 @@ func main() {
 	httpClient.Logger = importer.RetryableHTTPLeveledLogger{}
 	config.HTTPClient = httpClient.StandardClient()
 
-	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	ctx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
 	defer stop()
 
 	datastoreClient, err := datastore.NewClient(ctx, project)
 	if err != nil {
-		logger.Fatal("Failed to create datastore client", slog.Any("error", err))
+		logger.FatalContext(ctx, "Failed to create datastore client", slog.Any("error", err))
 	}
 	config.SourceRepoStore = db.NewSourceRepositoryStore(datastoreClient)
 	// Needed for deletions only
 	config.VulnerabilityStore = db.NewVulnerabilityStore(datastoreClient)
 
 	psClient, err := pubsub.NewClient(ctx, project)
 	if err != nil {
-		logger.Fatal("Failed to create pubsub client", slog.Any("error", err))
+		logger.FatalContext(ctx, "Failed to create pubsub client", slog.Any("error", err))
 	}
 	config.Publisher = &clients.GCPPublisher{Publisher: psClient.Publisher(importer.TasksTopic)}
 
-	storageClient, err := storage.NewClient(ctx)
+	// We are posssibly reading a lot of vulnerabilities from GCS, so disable telemetry (disables trace spans).
+	storageClient, err := storage.NewClient(ctx, option.WithTelemetryDisabled())
 	if err != nil {
-		logger.Fatal("Failed to create GCS client", slog.Any("error", err))
+		logger.FatalContext(ctx, "Failed to create GCS client", slog.Any("error", err))
 	}
 	config.GCSProvider = clients.NewGCSStorageProvider(storageClient)
 
 	if *runDelete {
 		if err := importer.RunDeletions(ctx, config); err != nil {
-			logger.Fatal("Importer-deleter failed", slog.Any("error", err))
+			logger.FatalContext(ctx, "Importer-deleter failed", slog.Any("error", err))
 		}
 	} else {
 		if err := importer.Run(ctx, config); err != nil {
-			logger.Fatal("Importer failed", slog.Any("error", err))
+			logger.FatalContext(ctx, "Importer failed", slog.Any("error", err))
 		}
 	}
 }
+
+// importerSampleRate returns the sample rate for the importer (not the individual vulnerability entries).
+// It is set to 0.05 (5%) by default, but can be overridden by the
+// IMPORT_TRACE_SAMPLE_RATE environment variable.
+func importerSampleRate() float64 {
+	rate := 0.05
+	if val := os.Getenv("IMPORT_TRACE_SAMPLE_RATE"); val != "" {
+		rate, _ = strconv.ParseFloat(val, 64)
+	}
+
+	return rate
+}
+
+// vulnerabilitySampleRate returns the sample rate for individual vulnerability entries.
+// It is set to 0.05 (5%) by default, but can be overridden by the
+// TRACE_SAMPLE_RATE environment variable.
+func vulnerabilitySampleRate() float64 {
+	rate := 0.05
+	if val := os.Getenv("TRACE_SAMPLE_RATE"); val != "" {
+		rate, _ = strconv.ParseFloat(val, 64)
+	}
+
+	return rate
+}

go/internal/importer/bucket.go

@@ -16,23 +16,13 @@ import (
 )
 
 type bucketSourceRecord struct {
-	bucket           clients.CloudStorage
-	objectPath       string
-	keyPath          string
-	hasUpdateTime    bool
-	lastUpdated      time.Time
-	format           RecordFormat
-	sourceRepository string
-	strict           bool
-	isDeleted        bool
+	bucket     clients.CloudStorage
+	objectPath string
 }
 
 var _ SourceRecord = bucketSourceRecord{}
 
 func (b bucketSourceRecord) Open(ctx context.Context) (io.ReadCloser, error) {
-	if b.isDeleted {
-		return nil, errors.New("cannot open a deleted record")
-	}
 	data, err := b.bucket.ReadObject(ctx, b.objectPath)
 	if err != nil {
 		return nil, err
@@ -41,43 +31,11 @@ func (b bucketSourceRecord) Open(ctx context.Context) (io.ReadCloser, error) {
 	return io.NopCloser(bytes.NewReader(data)), nil
 }
 
-func (b bucketSourceRecord) KeyPath() string {
-	return b.keyPath
-}
-
-func (b bucketSourceRecord) Format() RecordFormat {
-	return b.format
-}
-
-func (b bucketSourceRecord) LastUpdated() (time.Time, bool) {
-	return b.lastUpdated, b.hasUpdateTime
-}
-
-func (b bucketSourceRecord) SourceRepository() string {
-	return b.sourceRepository
-}
-
-func (b bucketSourceRecord) SourcePath() string {
-	return b.objectPath
-}
-
-func (b bucketSourceRecord) ShouldSendModifiedTime() bool {
-	return b.hasUpdateTime
-}
-
-func (b bucketSourceRecord) IsDeleted() bool {
-	return b.isDeleted
-}
-
-func (b bucketSourceRecord) Strictness() bool {
-	return b.strict
-}
-
-func handleImportBucket(ctx context.Context, ch chan<- SourceRecord, config Config, sourceRepo *models.SourceRepository) error {
+func handleImportBucket(ctx context.Context, ch chan<- WorkItem, config Config, sourceRepo *models.SourceRepository) error {
 	if sourceRepo.Type != models.SourceRepositoryTypeBucket || sourceRepo.Bucket == nil {
 		return errors.New("invalid SourceRepository for bucket import")
 	}
-	logger.Info("Importing bucket source repository",
+	logger.InfoContext(ctx, "Importing bucket source repository",
 		slog.String("source", sourceRepo.Name), slog.String("bucket", sourceRepo.Bucket.Name))
 
 	compiledIgnorePatterns := compileIgnorePatterns(sourceRepo)
@@ -100,7 +58,7 @@ func handleImportBucket(ctx context.Context, ch chan<- SourceRecord, config Conf
 			return err
 		}
 		if hasUpdateTime {
-			if obj.Attrs.CustomTime.Before(lastUpdated) {
+			if obj.Attrs.Updated.Before(lastUpdated) {
 				continue
 			}
 		}
@@ -111,37 +69,42 @@ func handleImportBucket(ctx context.Context, ch chan<- SourceRecord, config Conf
 		if shouldIgnore(base, sourceRepo.IDPrefixes, compiledIgnorePatterns) {
 			continue
 		}
-		ch <- bucketSourceRecord{
-			bucket:           bucket,
-			objectPath:       obj.Name,
-			lastUpdated:      lastUpdated,
-			hasUpdateTime:    hasUpdateTime,
-			format:           format,
-			keyPath:          sourceRepo.KeyPath,
-			sourceRepository: sourceRepo.Name,
-			strict:           sourceRepo.Strictness,
+		ch <- WorkItem{
+			Context: ctx,
+			SourceRecord: bucketSourceRecord{
+				bucket:     bucket,
+				objectPath: obj.Name,
+			},
+			SourceRepository:       sourceRepo.Name,
+			SourcePath:             obj.Name,
+			LastUpdated:            lastUpdated,
+			HasLastUpdated:         hasUpdateTime,
+			Format:                 format,
+			KeyPath:                sourceRepo.KeyPath,
+			Strict:                 sourceRepo.Strictness,
+			ShouldSendModifiedTime: hasUpdateTime,
 		}
 	}
 
 	sourceRepo.Bucket.LastUpdated = &timeOfRun
 	sourceRepo.Bucket.IgnoreLastImportTime = false
 	if err := config.SourceRepoStore.Update(ctx, sourceRepo.Name, sourceRepo); err != nil {
-		logger.Error("Failed to update source repository", slog.Any("error", err), slog.String("source", sourceRepo.Name))
+		logger.ErrorContext(ctx, "Failed to update source repository", slog.Any("error", err), slog.String("source", sourceRepo.Name))
 		return err
 	}
-	logger.Info("Finished importing bucket source repository",
+	logger.InfoContext(ctx, "Finished importing bucket source repository",
 		slog.String("source", sourceRepo.Name),
 		slog.String("bucket", sourceRepo.Bucket.Name))
 
 	return nil
 }
 
-func handleDeleteBucket(ctx context.Context, ch chan<- SourceRecord, config Config, sourceRepo *models.SourceRepository) error {
+func handleDeleteBucket(ctx context.Context, ch chan<- WorkItem, config Config, sourceRepo *models.SourceRepository) error {
 	if sourceRepo.Type != models.SourceRepositoryTypeBucket || sourceRepo.Bucket == nil {
 		return errors.New("invalid SourceRepository for bucket deletion")
 	}
 
-	logger.Info("Processing bucket deletions",
+	logger.InfoContext(ctx, "Processing bucket deletions",
 		slog.String("source", sourceRepo.Name), slog.String("bucket", sourceRepo.Bucket.Name))
 
 	// Get all objects in the bucket
@@ -167,7 +130,7 @@ func handleDeleteBucket(ctx context.Context, ch chan<- SourceRecord, config Conf
 	}
 
 	if len(vulnsInDatastore) == 0 {
-		logger.Info("No vulnerabilities found in Datastore for source", slog.String("source", sourceRepo.Name))
+		logger.InfoContext(ctx, "No vulnerabilities found in Datastore for source", slog.String("source", sourceRepo.Name))
 		return nil
 	}
 
@@ -180,7 +143,7 @@ func handleDeleteBucket(ctx context.Context, ch chan<- SourceRecord, config Conf
 	}
 
 	if len(toDelete) == 0 {
-		logger.Info("No vulnerabilities to delete", slog.String("source", sourceRepo.Name))
+		logger.InfoContext(ctx, "No vulnerabilities to delete", slog.String("source", sourceRepo.Name))
 		return nil
 	}
 
@@ -191,7 +154,7 @@ func handleDeleteBucket(ctx context.Context, ch chan<- SourceRecord, config Conf
 	}
 	percentage := (float64(len(toDelete)) / float64(len(vulnsInDatastore))) * 100.0
 	if percentage >= threshold {
-		logger.Error("Cowardly refusing to delete missing records (threshold exceeded)",
+		logger.ErrorContext(ctx, "Cowardly refusing to delete missing records (threshold exceeded)",
 			slog.String("source", sourceRepo.Name),
 			slog.Int("to_delete", len(toDelete)),
 			slog.Int("total", len(vulnsInDatastore)),
@@ -203,10 +166,15 @@ func handleDeleteBucket(ctx context.Context, ch chan<- SourceRecord, config Conf
 
 	// Trigger deletions
 	for _, entry := range toDelete {
-		ch <- bucketSourceRecord{
-			sourceRepository: entry.Source,
-			objectPath:       entry.Path,
-			isDeleted:        true,
+		ch <- WorkItem{
+			Context: ctx,
+			SourceRecord: bucketSourceRecord{
+				bucket:     bucket,
+				objectPath: entry.Path,
+			},
+			SourceRepository: entry.Source,
+			SourcePath:       entry.Path,
+			IsDeleted:        true,
 		}
 	}
 

go/internal/importer/bucket_test.go

@@ -48,19 +48,12 @@ func TestHandleImportBucket(t *testing.T) {
 
 	// Set up mock bucket with a file
 	mockBucket := testutils.NewMockStorage()
-	lastUpdated := time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC)
-	olderTime := time.Date(2022, 1, 1, 0, 0, 0, 0, time.UTC)
-	newerTime := time.Date(2023, 2, 1, 0, 0, 0, 0, time.UTC)
-
-	_ = mockBucket.WriteObject(ctx, "a/b/valid.json", []byte(`{}`), &clients.WriteOptions{
-		CustomTime: &olderTime, // older than last updated, will be skipped
-	})
-	_ = mockBucket.WriteObject(ctx, "a/b/newer.json", []byte(`{}`), &clients.WriteOptions{
-		CustomTime: &newerTime,
-	})
-	_ = mockBucket.WriteObject(ctx, "a/b/ignored.json", []byte(`{}`), &clients.WriteOptions{
-		CustomTime: &newerTime,
-	})
+	_ = mockBucket.WriteObject(ctx, "a/b/valid.json", []byte(`{}`), &clients.WriteOptions{})
+	time.Sleep(50 * time.Millisecond)
+	lastUpdated := time.Now()
+	time.Sleep(50 * time.Millisecond)
+	_ = mockBucket.WriteObject(ctx, "a/b/newer.json", []byte(`{}`), &clients.WriteOptions{})
+	_ = mockBucket.WriteObject(ctx, "a/b/ignored.json", []byte(`{}`), &clients.WriteOptions{})
 
 	provider := &mockCloudStorageProvider{
 		buckets: map[string]clients.CloudStorage{
@@ -89,7 +82,7 @@ func TestHandleImportBucket(t *testing.T) {
 		},
 	}
 
-	ch := make(chan SourceRecord, 10)
+	ch := make(chan WorkItem, 10)
 
 	err := handleImportBucket(ctx, ch, config, sourceRepo)
 	if err != nil {
@@ -99,7 +92,7 @@ func TestHandleImportBucket(t *testing.T) {
 	close(ch)
 	records := make([]bucketSourceRecord, 0, 10)
 	for r := range ch {
-		records = append(records, r.(bucketSourceRecord))
+		records = append(records, r.SourceRecord.(bucketSourceRecord))
 	}
 
 	if len(records) != 1 {