Bound the direction precompute shape cache

gtfsdb/query.sql

@@ -415,12 +415,6 @@ FROM
 WHERE
     stop_times.stop_id = ?;
 
--- name: GetAllShapes :many
-SELECT
-    *
-FROM
-    shapes;
-
 -- name: GetShapeByID :many
 SELECT
     *

internal/gtfs/advanced_direction_calculator.go

@@ -26,10 +26,9 @@ const (
 type AdvancedDirectionCalculator struct {
 	queries                    *gtfsdb.Queries
 	standardDeviationThreshold float64
-	shapeCache                 map[string][]gtfsdb.GetShapePointsWithDistanceRow // Cache of all shape data for bulk operations
-	initialized                atomic.Bool                                       // Tracks whether concurrent operations have started
-	cacheMutex                 sync.RWMutex                                      // Protects shapeCache map access
 	cache                      atomic.Pointer[directionCache]
+	recentShapes               *recentShapes      // optional; only for a calculator that lives for one import, since ClearCache doesn't empty it
+	shapeLoads                 singleflight.Group // merges concurrent reads of one shape into recentShapes
 }
 
 // Each calculation retains its generation so work in flight during a reload cannot
@@ -57,22 +56,6 @@ func (adc *AdvancedDirectionCalculator) ClearCache() {
 	adc.cache.Store(&directionCache{})
 }
 
-// SetShapeCache is retained exclusively for use by the DirectionPrecomputer during startup.
-// It sets a pre-loaded cache of shape data to avoid thousands of database queries during
-// the precomputation phase, significantly improving startup performance.
-// IMPORTANT: This must be called before any concurrent operations begin.
-// Returns an error if called after CalculateStopDirection has been invoked.
-func (adc *AdvancedDirectionCalculator) SetShapeCache(cache map[string][]gtfsdb.GetShapePointsWithDistanceRow) error {
-	adc.cacheMutex.Lock()
-	defer adc.cacheMutex.Unlock()
-
-	if adc.initialized.Load() {
-		return errors.New("SetShapeCache called after concurrent operations have started")
-	}
-	adc.shapeCache = cache
-	return nil
-}
-
 // shapeDistancesCanLocateAStop reports whether a shape's shape_dist_traveled values
 // are usable for finding where along it a stop sits.
 //
@@ -160,9 +143,6 @@ func (adc *AdvancedDirectionCalculator) CalculateStopDirection(ctx context.Conte
 		return cached.(string)
 	}
 
-	// Mark as initialized for concurrency safety
-	adc.initialized.Store(true)
-
 	// Fall back to computing from shapes, protected by singleflight
 	// This ensures concurrent requests for the SAME stopID don't hit the DB multiple times.
 	v, _, _ := cache.requestGroup.Do(stopID, func() (any, error) {
@@ -372,36 +352,45 @@ func (adc *AdvancedDirectionCalculator) directionFromOrientations(orientations [
 	return adc.getAngleAsDirection(thetaMedian)
 }
 
-// calculateOrientationAtStop calculates the orientation at a stop using a window of shape points
-// If distTraveled is < 0, it uses stopLat/stopLon for geographic matching (when shape_dist_traveled is unavailable)
-func (adc *AdvancedDirectionCalculator) calculateOrientationAtStop(ctx context.Context, shapeID string, distTraveled float64, stopLat, stopLon float64) (float64, error) {
-	var shapePoints []gtfsdb.GetShapePointsWithDistanceRow
-	var err error
-
-	adc.cacheMutex.RLock()
-	hasCache := adc.shapeCache != nil
-	var found bool
-	if hasCache {
-		shapePoints, found = adc.shapeCache[shapeID]
+// shapePoints reads a shape's points from the database, through recentShapes
+// when the calculator has one. Cache hits return before singleflight, which
+// only merges concurrent reads of a shape that isn't cached.
+func (adc *AdvancedDirectionCalculator) shapePoints(ctx context.Context, shapeID string) ([]gtfsdb.GetShapePointsWithDistanceRow, error) {
+	if adc.recentShapes == nil {
+		return adc.queries.GetShapePointsWithDistance(ctx, shapeID)
 	}
-	adc.cacheMutex.RUnlock()
-
-	// Try cache first if available
-	if hasCache {
-		if !found || len(shapePoints) < 2 {
-			return 0, sql.ErrNoRows
+	if points, ok := adc.recentShapes.get(shapeID); ok {
+		return points, nil
+	}
+	loaded, err, _ := adc.shapeLoads.Do(shapeID, func() (any, error) {
+		// A read that finished after the get above has already cached the shape.
+		if points, ok := adc.recentShapes.peek(shapeID); ok {
+			return points, nil
 		}
-	} else {
-		// Fall back to database query if no cache
-		shapePoints, err = adc.queries.GetShapePointsWithDistance(ctx, shapeID)
+		points, err := adc.queries.GetShapePointsWithDistance(ctx, shapeID)
 		if err != nil {
-			return 0, err
-		}
-		if len(shapePoints) < 2 {
-			// Insufficient points is a data condition, not a transient error.
-			// Return ErrNoRows so the caller treats this the same as "no shape data".
-			return 0, sql.ErrNoRows
+			return nil, err
 		}
+		adc.recentShapes.add(shapeID, points)
+		return points, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return loaded.([]gtfsdb.GetShapePointsWithDistanceRow), nil
+}
+
+// calculateOrientationAtStop calculates the orientation at a stop using a window of shape points
+// If distTraveled is < 0, it uses stopLat/stopLon for geographic matching (when shape_dist_traveled is unavailable)
+func (adc *AdvancedDirectionCalculator) calculateOrientationAtStop(ctx context.Context, shapeID string, distTraveled float64, stopLat, stopLon float64) (float64, error) {
+	shapePoints, err := adc.shapePoints(ctx, shapeID)
+	if err != nil {
+		return 0, err
+	}
+	if len(shapePoints) < 2 {
+		// Insufficient points is a data condition, not a transient error.
+		// Return ErrNoRows so the caller treats this the same as "no shape data".
+		return 0, sql.ErrNoRows
 	}
 
 	closestIdx := 0

internal/gtfs/advanced_direction_calculator_test.go

@@ -557,8 +557,9 @@ func shapeRow(lat, lon float64, seq int64) gtfsdb.GetShapePointsWithDistanceRow
 func newCachedCalc(t *testing.T, shapes map[string][]gtfsdb.GetShapePointsWithDistanceRow) *AdvancedDirectionCalculator {
 	t.Helper()
 	calc := NewAdvancedDirectionCalculator(nil)
-	if err := calc.SetShapeCache(shapes); err != nil {
-		t.Fatalf("SetShapeCache: %v", err)
+	calc.recentShapes = newRecentShapes(math.MaxInt)
+	for shapeID, points := range shapes {
+		calc.recentShapes.add(shapeID, points)
 	}
 	return calc
 }

internal/gtfs/direction_precomputer.go

@@ -19,6 +19,11 @@ type stopDirectionResult struct {
 	direction string
 }
 
+// precomputeShapeCachePoints is how many shape points the precomputer keeps
+// in memory, about 20 MB of rows. It holds roughly 500 shapes of average size
+// in a large bus feed such as NJ Transit's.
+const precomputeShapeCachePoints = 500_000
+
 // DirectionPrecomputer handles batch computation of stop directions
 type DirectionPrecomputer struct {
 	calculator *AdvancedDirectionCalculator
@@ -29,8 +34,10 @@ type DirectionPrecomputer struct {
 
 // NewDirectionPrecomputer creates a new direction precomputer
 func NewDirectionPrecomputer(queries *gtfsdb.Queries, db *sql.DB) *DirectionPrecomputer {
+	calculator := NewAdvancedDirectionCalculator(queries)
+	calculator.recentShapes = newRecentShapes(precomputeShapeCachePoints)
 	return &DirectionPrecomputer{
-		calculator: NewAdvancedDirectionCalculator(queries),
+		calculator: calculator,
 		queries:    queries,
 		db:         db,
 		logger:     slog.Default().With(slog.String("component", "direction_precomputer")),
@@ -54,35 +61,6 @@ func (dp *DirectionPrecomputer) PrecomputeAllDirections(ctx context.Context) err
 		return nil
 	}
 
-	// ===== PHASE 0: PRE-LOAD SHAPE DATA INTO CACHE =====
-	logging.LogOperation(dp.logger, "loading_shape_cache_started")
-	shapeCache, err := dp.loadShapeCache(ctx)
-	if err != nil {
-		logging.LogError(dp.logger, "Failed to load shape cache", err)
-		// Don't fail precomputation if cache loading fails - will fall back to DB queries
-	} else {
-		// Calculate total points across all shapes
-		totalPoints := 0
-		for _, points := range shapeCache {
-			totalPoints += len(points)
-		}
-
-		// Estimate memory usage (approximate):
-		// Each GetShapePointsWithDistanceRow: 2 float64 (lat/lon) + 1 int64 (sequence) + 1 sql.NullFloat64 (dist)
-		// = 16 + 16 + 8 + 24 = 64 bytes per point (conservative estimate)
-		// Plus map overhead and slice headers
-		estimatedMemoryMB := float64(totalPoints*64) / (1024 * 1024)
-
-		logging.LogOperation(dp.logger, "shape_cache_loaded",
-			slog.Int("shape_count", len(shapeCache)),
-			slog.Int("total_points", totalPoints),
-			slog.Float64("estimated_memory_mb", estimatedMemoryMB))
-		// Set the cache on the calculator for use during precomputation
-		if err := dp.calculator.SetShapeCache(shapeCache); err != nil {
-			return fmt.Errorf("failed to set shape cache: %w", err)
-		}
-	}
-
 	// ===== PHASE 1: PARALLEL DIRECTION CALCULATION (read-only) =====
 	numWorkers := runtime.NumCPU()
 	stopsChan := make(chan gtfsdb.Stop, numWorkers)
@@ -236,39 +214,16 @@ func (dp *DirectionPrecomputer) PrecomputeAllDirections(ctx context.Context) err
 	}
 
 	duration := time.Since(startTime)
+	shapeStats := dp.calculator.recentShapes.lookupStats()
 	logging.LogOperation(dp.logger, "precomputing_stop_directions_completed",
 		slog.Duration("duration_ms", duration),
 		slog.Int("total_stops", len(stops)),
 		slog.Int("successful", successCount),
 		slog.Int("skipped", skippedCount),
-		slog.Int("errors", errorCount))
+		slog.Int("errors", errorCount),
+		slog.Int("shape_lookups", shapeStats.lookups),
+		slog.Int("shape_cache_hits", shapeStats.hits),
+		slog.Int("shape_reads", shapeStats.reads))
 
 	return nil
 }
-
-// loadShapeCache loads all shape data from the database and organizes it by shape_id
-// for efficient lookup during direction precomputation
-func (dp *DirectionPrecomputer) loadShapeCache(ctx context.Context) (map[string][]gtfsdb.GetShapePointsWithDistanceRow, error) {
-	// Get all shape points from the database
-	allShapes, err := dp.queries.GetAllShapes(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("failed to fetch all shapes: %w", err)
-	}
-
-	// Organize shapes by shape_id for O(1) lookup
-	shapeCache := make(map[string][]gtfsdb.GetShapePointsWithDistanceRow)
-
-	for _, shape := range allShapes {
-		// Convert Shape to GetShapePointsWithDistanceRow (excludes ShapeID as it's the map key)
-		shapePoint := gtfsdb.GetShapePointsWithDistanceRow{
-			Lat:               shape.Lat,
-			Lon:               shape.Lon,
-			ShapePtSequence:   shape.ShapePtSequence,
-			ShapeDistTraveled: shape.ShapeDistTraveled,
-		}
-
-		shapeCache[shape.ShapeID] = append(shapeCache[shape.ShapeID], shapePoint)
-	}
-
-	return shapeCache, nil
-}