feat: Add callback support in code graph builder

code.go

@@ -86,7 +86,7 @@ func internalStartImportReachability() error {
 		return fmt.Errorf("no database path provided")
 	}
 
-	// TODO: We need a CPG loader to load the CPG from the database
+	// TODO: We need a code graph loader to load the code graph from the database
 	// before invoking analysis modules
 
 	return nil
@@ -140,14 +140,46 @@ func internalStartCreateDatabase() error {
 
 	builder, err := code.NewCodeGraphBuilder(builderConfig, codeRepo, codeLang, graph)
 	if err != nil {
-		return fmt.Errorf("failed to create CPG builder: %w", err)
+		return fmt.Errorf("failed to create code graph builder: %w", err)
 	}
 
+	redirectLogToFile(logFile)
+
+	var fileProcessedTracker any
+	var importsProcessedTracker any
+	var functionsProcessedTracker any
+
+	builder.RegisterEventHandler("ui-callback",
+		func(event code.CodeGraphBuilderEvent, metrics code.CodeGraphBuilderMetrics) error {
+			switch event.Kind {
+			case code.CodeGraphBuilderEventFileQueued:
+				ui.IncrementTrackerTotal(fileProcessedTracker, 1)
+			case code.CodeGraphBuilderEventFileProcessed:
+				ui.IncrementProgress(fileProcessedTracker, 1)
+			}
+
+			ui.UpdateValue(importsProcessedTracker, int64(metrics.ImportsCount))
+			ui.UpdateValue(functionsProcessedTracker, int64(metrics.FunctionsCount))
+
+			return nil
+		})
+
+	ui.StartProgressWriter()
+
+	fileProcessedTracker = ui.TrackProgress("Processing source files", 0)
+	importsProcessedTracker = ui.TrackProgress("Processing imports", 0)
+	functionsProcessedTracker = ui.TrackProgress("Processing functions", 0)
+
 	err = builder.Build()
 	if err != nil {
-		return fmt.Errorf("failed to build CPG: %w", err)
+		return fmt.Errorf("failed to build code graph: %w", err)
 	}
 
+	ui.MarkTrackerAsDone(fileProcessedTracker)
+	ui.MarkTrackerAsDone(importsProcessedTracker)
+	ui.MarkTrackerAsDone(functionsProcessedTracker)
+	ui.StopProgressWriter()
+
 	logger.Debugf("Code analysis completed")
 	return nil
 }

internal/ui/progress.go

@@ -14,7 +14,7 @@ func StartProgressWriter() {
 
 	pw.SetAutoStop(false)
 	pw.SetTrackerLength(25)
-	pw.SetMessageWidth(20)
+	pw.SetMessageLength(20)
 	pw.SetSortBy(progress.SortByPercentDsc)
 	pw.SetStyle(progress.StyleDefault)
 	pw.SetOutputWriter(os.Stderr)
@@ -72,6 +72,12 @@ func IncrementProgress(i any, count int64) {
 	}
 }
 
+func UpdateValue(i any, count int64) {
+	if tracker, ok := i.(*progress.Tracker); ok {
+		tracker.SetValue(count)
+	}
+}
+
 func progressTrackerDelta(tracker *progress.Tracker) int64 {
 	return (tracker.Total - tracker.Value())
 }

pkg/code/code_graph.go

@@ -17,8 +17,35 @@ type CodeGraphBuilderConfig struct {
 	Concurrency int
 }
 
+type CodeGraphBuilderMetrics struct {
+	GoRoutineCount int
+	ErrorCount     int
+	FilesProcessed int
+	FilesInQueue   int
+	ImportsCount   int
+	FunctionsCount int
+}
+
+type CodeGraphBuilderEvent struct {
+	Kind string
+	Data interface{}
+}
+
+const (
+	CodeGraphBuilderEventFileQueued        = "file_queued"
+	CodeGraphBuilderEventFileProcessed     = "file_processed"
+	CodeGraphBuilderEventImportProcessed   = "import_processed"
+	CodeGraphBuilderEventFunctionProcessed = "function_processed"
+)
+
+// The event handler is invoked from its own go routine. Handler functions
+// must be thread safe.
+type CodeGraphBuilderEventHandler func(CodeGraphBuilderEvent, CodeGraphBuilderMetrics) error
+
 type codeGraphBuilder struct {
-	config CodeGraphBuilderConfig
+	config        CodeGraphBuilderConfig
+	eventHandlers map[string]CodeGraphBuilderEventHandler
+	metrics       CodeGraphBuilderMetrics
 
 	repository SourceRepository
 	lang       SourceLanguage
@@ -44,9 +71,21 @@ func NewCodeGraphBuilder(config CodeGraphBuilderConfig,
 	}
 
 	return &codeGraphBuilder{config: config,
-		repository: repository,
-		lang:       lang,
-		storage:    storage}, nil
+		repository:    repository,
+		lang:          lang,
+		storage:       storage,
+		eventHandlers: make(map[string]CodeGraphBuilderEventHandler, 0),
+		metrics:       CodeGraphBuilderMetrics{GoRoutineCount: config.Concurrency},
+	}, nil
+}
+
+func (b *codeGraphBuilder) RegisterEventHandler(name string, handler CodeGraphBuilderEventHandler) {
+	if _, ok := b.eventHandlers[name]; ok {
+		logger.Warnf("Event handler already registered: %s", name)
+		return
+	}
+
+	b.eventHandlers[name] = handler
 }
 
 func (b *codeGraphBuilder) Build() error {
@@ -83,29 +122,48 @@ func (b *codeGraphBuilder) Build() error {
 }
 
 func (b *codeGraphBuilder) enqueueSourceFile(file SourceFile) {
-	b.fileQueueLock.Lock()
-	defer b.fileQueueLock.Unlock()
+	b.synchronized(func() {
+		b.metrics.FilesInQueue++
 
-	if _, ok := b.fileCache[file.Path]; ok {
-		logger.Debugf("Skipping already processed file: %s", file.Path)
-		return
-	}
+		if _, ok := b.fileCache[file.Path]; ok {
+			logger.Debugf("Skipping already processed file: %s", file.Path)
+			return
+		}
+
+		b.fileQueueWg.Add(1)
+		b.fileQueue <- file
 
-	b.fileQueueWg.Add(1)
-	b.fileQueue <- file
+		// TODO: Optimize this. Storing entire file path is not needed
+		b.fileCache[file.Path] = true
+	})
 
-	// TODO: Optimize this. Storing entire file path is not needed
-	b.fileCache[file.Path] = true
+	b.notifyEventHandlers(CodeGraphBuilderEvent{
+		Kind: CodeGraphBuilderEventFileQueued,
+		Data: file,
+	}, b.metrics)
 }
 
 func (b *codeGraphBuilder) fileProcessor(wg *sync.WaitGroup) {
 	for file := range b.fileQueue {
 		err := b.buildForFile(file)
 		if err != nil {
 			logger.Errorf("Failed to process code graph for: %s: %v", file.Path, err)
+
+			b.synchronized(func() {
+				b.metrics.ErrorCount++
+			})
 		}
 
 		wg.Done()
+
+		b.synchronized(func() {
+			b.metrics.FilesProcessed++
+		})
+
+		b.notifyEventHandlers(CodeGraphBuilderEvent{
+			Kind: CodeGraphBuilderEventFileProcessed,
+			Data: file,
+		}, b.metrics)
 	}
 }
 
@@ -187,6 +245,15 @@ func (b *codeGraphBuilder) processImportNodes(cst *nodes.CST, currentFile Source
 		if err != nil {
 			logger.Errorf("Failed to link import node: %v", err)
 		}
+
+		b.synchronized(func() {
+			b.metrics.ImportsCount++
+		})
+
+		b.notifyEventHandlers(CodeGraphBuilderEvent{
+			Kind: CodeGraphBuilderEventImportProcessed,
+			Data: importNode.ImportName(),
+		}, b.metrics)
 	}
 }
 
@@ -230,6 +297,15 @@ func (b *codeGraphBuilder) processFunctionDeclarations(cst *nodes.CST, currentFi
 			logger.Errorf("Failed to link function declaration: %v", err)
 		}
 
+		b.synchronized(func() {
+			b.metrics.FunctionsCount++
+		})
+
+		b.notifyEventHandlers(CodeGraphBuilderEvent{
+			Kind: CodeGraphBuilderEventFunctionProcessed,
+			Data: functionDecl.Id(),
+		}, b.metrics)
+
 		b.functionDeclCache[moduleName] = functionDecl.Id()
 	}
 }
@@ -269,3 +345,19 @@ func (b *codeGraphBuilder) importSourceName(file SourceFile) string {
 		return entities.PackageEntitySourceTypeApp
 	}
 }
+
+func (b *codeGraphBuilder) notifyEventHandlers(event CodeGraphBuilderEvent, metrics CodeGraphBuilderMetrics) {
+	for name, handler := range b.eventHandlers {
+		err := handler(event, metrics)
+		if err != nil {
+			logger.Warnf("Failed to notify event handler: %s: %v", name, err)
+		}
+	}
+}
+
+func (b *codeGraphBuilder) synchronized(fn func()) {
+	b.fileQueueLock.Lock()
+	defer b.fileQueueLock.Unlock()
+
+	fn()
+}