package storage import ( "context" "fmt" "sync" "time" "github.com/anthonyrawlins/bzzz/pkg/ucxl" slurpContext "github.com/anthonyrawlins/bzzz/pkg/slurp/context" ) // BatchOperationsImpl provides efficient batch operations for context storage type BatchOperationsImpl struct { contextStore *ContextStoreImpl batchSize int maxConcurrency int operationTimeout time.Duration } // NewBatchOperations creates a new batch operations handler func NewBatchOperations(contextStore *ContextStoreImpl, batchSize, maxConcurrency int, timeout time.Duration) *BatchOperationsImpl { return &BatchOperationsImpl{ contextStore: contextStore, batchSize: batchSize, maxConcurrency: maxConcurrency, operationTimeout: timeout, } } // BatchStore stores multiple contexts efficiently func (cs *ContextStoreImpl) BatchStore( ctx context.Context, batch *BatchStoreRequest, ) (*BatchStoreResult, error) { start := time.Now() result := &BatchStoreResult{ Errors: make(map[string]error), ProcessedAt: time.Now(), } // Validate batch request if batch == nil || len(batch.Contexts) == 0 { return result, fmt.Errorf("empty batch request") } // Create worker pool for concurrent processing workerCount := cs.options.MaxConcurrentOps if len(batch.Contexts) < workerCount { workerCount = len(batch.Contexts) } // Channels for work distribution workCh := make(chan *BatchStoreWork, len(batch.Contexts)) resultsCh := make(chan *BatchStoreWorkResult, len(batch.Contexts)) // Start workers var wg sync.WaitGroup for i := 0; i < workerCount; i++ { wg.Add(1) go func() { defer wg.Done() cs.batchStoreWorker(ctx, workCh, resultsCh, batch) }() } // Send work to workers go func() { defer close(workCh) for i, contextItem := range batch.Contexts { work := &BatchStoreWork{ Index: i, Item: contextItem, Timeout: cs.options.OperationTimeout, } workCh <- work } }() // Collect results go func() { wg.Wait() close(resultsCh) }() // Process results for workResult := range resultsCh { if workResult.Error != nil { result.ErrorCount++ key := workResult.Item.Context.UCXLAddress.String() result.Errors[key] = workResult.Error if batch.FailOnError { // Cancel remaining operations result.ProcessingTime = time.Since(start) return result, fmt.Errorf("batch operation failed on context %s: %w", key, workResult.Error) } } else { result.SuccessCount++ } } result.ProcessingTime = time.Since(start) return result, nil } // BatchRetrieve retrieves multiple contexts efficiently func (cs *ContextStoreImpl) BatchRetrieve( ctx context.Context, batch *BatchRetrieveRequest, ) (*BatchRetrieveResult, error) { start := time.Now() result := &BatchRetrieveResult{ Contexts: make(map[string]*slurpContext.ContextNode), Errors: make(map[string]error), ProcessedAt: time.Now(), } // Validate batch request if batch == nil || len(batch.Addresses) == 0 { return result, fmt.Errorf("empty batch request") } // Create worker pool for concurrent processing workerCount := cs.options.MaxConcurrentOps if len(batch.Addresses) < workerCount { workerCount = len(batch.Addresses) } // Channels for work distribution workCh := make(chan *BatchRetrieveWork, len(batch.Addresses)) resultsCh := make(chan *BatchRetrieveWorkResult, len(batch.Addresses)) // Start workers var wg sync.WaitGroup for i := 0; i < workerCount; i++ { wg.Add(1) go func() { defer wg.Done() cs.batchRetrieveWorker(ctx, workCh, resultsCh, batch) }() } // Send work to workers go func() { defer close(workCh) for i, address := range batch.Addresses { work := &BatchRetrieveWork{ Index: i, Address: address, Role: batch.Role, Timeout: cs.options.OperationTimeout, } workCh <- work } }() // Collect results go func() { wg.Wait() close(resultsCh) }() // Process results for workResult := range resultsCh { addressStr := workResult.Address.String() if workResult.Error != nil { result.ErrorCount++ result.Errors[addressStr] = workResult.Error if batch.FailOnError { // Cancel remaining operations result.ProcessingTime = time.Since(start) return result, fmt.Errorf("batch operation failed on address %s: %w", addressStr, workResult.Error) } } else { result.SuccessCount++ result.Contexts[addressStr] = workResult.Context } } result.ProcessingTime = time.Since(start) return result, nil } // Worker functions and supporting types type BatchStoreWork struct { Index int Item *ContextStoreItem Timeout time.Duration } type BatchStoreWorkResult struct { Index int Item *ContextStoreItem Error error } type BatchRetrieveWork struct { Index int Address ucxl.Address Role string Timeout time.Duration } type BatchRetrieveWorkResult struct { Index int Address ucxl.Address Context *slurpContext.ContextNode Error error } // batchStoreWorker processes batch store work items func (cs *ContextStoreImpl) batchStoreWorker( ctx context.Context, workCh <-chan *BatchStoreWork, resultsCh chan<- *BatchStoreWorkResult, batch *BatchStoreRequest, ) { for work := range workCh { result := &BatchStoreWorkResult{ Index: work.Index, Item: work.Item, } // Create timeout context for this operation workCtx, cancel := context.WithTimeout(ctx, work.Timeout) defer cancel() // Determine roles to use roles := work.Item.Roles if len(roles) == 0 { roles = batch.Roles // Use batch default roles } // Perform the store operation if batch.Transaction { // Use transaction if requested result.Error = cs.storeContextWithTransaction(workCtx, work.Item.Context, roles) } else { // Regular store result.Error = cs.StoreContext(workCtx, work.Item.Context, roles) } resultsCh <- result } } // batchRetrieveWorker processes batch retrieve work items func (cs *ContextStoreImpl) batchRetrieveWorker( ctx context.Context, workCh <-chan *BatchRetrieveWork, resultsCh chan<- *BatchRetrieveWorkResult, batch *BatchRetrieveRequest, ) { for work := range workCh { result := &BatchRetrieveWorkResult{ Index: work.Index, Address: work.Address, } // Create timeout context for this operation workCtx, cancel := context.WithTimeout(ctx, work.Timeout) defer cancel() // Perform the retrieve operation contextNode, err := cs.RetrieveContext(workCtx, work.Address, work.Role) result.Context = contextNode result.Error = err resultsCh <- result } } // storeContextWithTransaction performs a store operation within a transaction context func (cs *ContextStoreImpl) storeContextWithTransaction( ctx context.Context, node *slurpContext.ContextNode, roles []string, ) error { // This would integrate with a transaction manager if available // For now, it's the same as regular store but could be enhanced // with rollback capabilities return cs.StoreContext(ctx, node, roles) } // ListContexts lists contexts matching criteria with optimized querying func (cs *ContextStoreImpl) ListContexts( ctx context.Context, criteria *ListCriteria, ) ([]*slurpContext.ContextNode, error) { start := time.Now() defer func() { cs.recordLatency("list", time.Since(start)) }() // Use search index if available and appropriate if cs.options.IndexingEnabled && cs.shouldUseSearchIndex(criteria) { return cs.listContextsViaSearch(ctx, criteria) } // Fallback to storage enumeration return cs.listContextsViaStorage(ctx, criteria) } // SearchContexts searches contexts using query criteria with advanced features func (cs *ContextStoreImpl) SearchContexts( ctx context.Context, query *SearchQuery, ) (*SearchResults, error) { start := time.Now() defer func() { cs.recordLatency("search", time.Since(start)) }() // Validate search query if query == nil { return nil, fmt.Errorf("search query cannot be nil") } // Use primary search index indexName := "primary" if cs.indexManager != nil { indexes, err := cs.indexManager.ListIndexes(ctx) if err == nil && len(indexes) > 0 { indexName = indexes[0] // Use first available index } } // Perform search results, err := cs.indexManager.Search(ctx, indexName, query) if err != nil { return nil, fmt.Errorf("search failed: %w", err) } // Post-process results for role-based filtering if len(query.Roles) > 0 { results = cs.filterResultsByRole(ctx, results, query.Roles) } // Apply additional filters that couldn't be done at index level results = cs.applyPostSearchFilters(ctx, results, query) cs.recordOperation("search") return results, nil } // Helper methods for optimized listing and searching func (cs *ContextStoreImpl) shouldUseSearchIndex(criteria *ListCriteria) bool { // Use search index if we have complex criteria that would benefit from indexing return len(criteria.Tags) > 0 || len(criteria.Technologies) > 0 || criteria.PathPattern != "" || criteria.MinConfidence > 0 } func (cs *ContextStoreImpl) listContextsViaSearch( ctx context.Context, criteria *ListCriteria, ) ([]*slurpContext.ContextNode, error) { // Convert list criteria to search query query := &SearchQuery{ Tags: criteria.Tags, Technologies: criteria.Technologies, Roles: criteria.Roles, MinConfidence: criteria.MinConfidence, Limit: criteria.Limit, Offset: criteria.Offset, SortBy: criteria.SortBy, SortOrder: criteria.SortOrder, IncludeStale: criteria.IncludeStale, } // Add path pattern as scope if provided if criteria.PathPattern != "" { query.Scope = []string{criteria.PathPattern} } // Perform search searchResults, err := cs.SearchContexts(ctx, query) if err != nil { return nil, err } // Extract contexts from search results contexts := make([]*slurpContext.ContextNode, len(searchResults.Results)) for i, result := range searchResults.Results { contexts[i] = result.Context } return contexts, nil } func (cs *ContextStoreImpl) listContextsViaStorage( ctx context.Context, criteria *ListCriteria, ) ([]*slurpContext.ContextNode, error) { // This would enumerate storage and apply filters // This is a simplified implementation - in practice, this would be more sophisticated var contexts []*slurpContext.ContextNode // For now, return empty list as this would require storage enumeration // In a real implementation, this would iterate through storage keys // and load contexts that match the criteria return contexts, nil } func (cs *ContextStoreImpl) filterResultsByRole( ctx context.Context, results *SearchResults, roles []string, ) *SearchResults { // Filter search results based on role access // This ensures users only see contexts they have access to filteredResults := make([]*SearchResult, 0, len(results.Results)) for _, result := range results.Results { // Check if any of the requested roles can access this context hasAccess := false for _, role := range roles { if cs.options.EncryptionEnabled { storageKey := cs.generateStorageKey(result.Context.UCXLAddress) if canAccess, err := cs.encryptedStorage.CanAccess(ctx, storageKey, role); err == nil && canAccess { hasAccess = true break } } else { // For unencrypted storage, assume access is allowed hasAccess = true break } } if hasAccess { filteredResults = append(filteredResults, result) } } // Update result metadata results.Results = filteredResults results.TotalResults = int64(len(filteredResults)) return results } func (cs *ContextStoreImpl) applyPostSearchFilters( ctx context.Context, results *SearchResults, query *SearchQuery, ) *SearchResults { // Apply filters that couldn't be applied at the search index level filteredResults := make([]*SearchResult, 0, len(results.Results)) for _, result := range results.Results { include := true // Age filter if query.MaxAge != nil { age := time.Since(result.Context.GeneratedAt) if age > *query.MaxAge { include = false } } // File type filter (based on path extension) if len(query.FileTypes) > 0 { matchesFileType := false for _, fileType := range query.FileTypes { if strings.HasSuffix(result.Context.Path, "."+fileType) { matchesFileType = true break } } if !matchesFileType { include = false } } // Scope filter if len(query.Scope) > 0 { matchesScope := false for _, scope := range query.Scope { if strings.HasPrefix(result.Context.Path, scope) { matchesScope = true break } } if !matchesScope { include = false } } // Exclude scope filter if len(query.ExcludeScope) > 0 { for _, excludeScope := range query.ExcludeScope { if strings.HasPrefix(result.Context.Path, excludeScope) { include = false break } } } if include { filteredResults = append(filteredResults, result) } } // Update result metadata results.Results = filteredResults results.TotalResults = int64(len(filteredResults)) return results }