<commit_before>package libkbfs

import (
	"sort"

	"golang.org/x/net/context"
)

const (
	defaultIndirectPointerPrefetchCount int = 20
	fileIndirectBlockPrefetchPriority   int = -100
	dirEntryPrefetchPriority            int = -200
)

type dirEntries []DirEntry
type dirEntriesBySizeAsc struct{ dirEntries }
type dirEntriesBySizeDesc struct{ dirEntries }

func (d dirEntries) Len() int                     { return len(d) }
func (d dirEntries) Swap(i, j int)                { d[i], d[j] = d[j], d[i] }
func (d dirEntriesBySizeAsc) Less(i, j int) bool  { return d.dirEntries[i].Size < d.dirEntries[j].Size }
func (d dirEntriesBySizeDesc) Less(i, j int) bool { return d.dirEntries[i].Size > d.dirEntries[j].Size }

func dirEntryMapToDirEntries(entryMap map[string]DirEntry) dirEntries {
	dirEntries := make(dirEntries, 0, len(entryMap))
	for _, entry := range entryMap {
		dirEntries = append(dirEntries, entry)
	}
	return dirEntries
}

type prefetcher interface {
	HandleBlock(b Block, kmd KeyMetadata, priority int) error
}

var _ prefetcher = (*blockPrefetcher)(nil)

type blockPrefetcher struct {
	retriever blockRetriever
}

func newPrefetcher(retriever blockRetriever) prefetcher {
	return &blockPrefetcher{
		retriever: retriever,
	}
}

func (p *blockPrefetcher) HandleBlock(b Block, kmd KeyMetadata, priority int) error {
	switch b := b.(type) {
	case *FileBlock:
		// If this is an indirect block and the priority is on demand, prefetch
		// the first <n> indirect block pointers.
		// TODO: do something smart with subsequent blocks.
		if b.IsInd && priority >= defaultOnDemandRequestPriority {
			numIPtrs := len(b.IPtrs)
			if numIPtrs > defaultIndirectPointerPrefetchCount {
				numIPtrs = defaultIndirectPointerPrefetchCount
			}
			for _, ptr := range b.IPtrs[:numIPtrs] {
				p.request(fileIndirectBlockPrefetchPriority, kmd,
					ptr.BlockPointer, NewFileBlock(), TransientEntry)
			}
		}
	case *DirBlock:
		// If this is an on-demand request:
		// - If the block is indirect, prefetch the first <n> indirect block
		// pointers.
		// - If the block is direct (has Children), prefetch all DirEntry root
		// blocks.
		if priority >= defaultOnDemandRequestPriority {
			if b.IsInd {
				numIPtrs := len(b.IPtrs)
				if numIPtrs > defaultIndirectPointerPrefetchCount {
					numIPtrs = defaultIndirectPointerPrefetchCount
				}
				for _, ptr := range b.IPtrs[:numIPtrs] {
					p.request(fileIndirectBlockPrefetchPriority, kmd,
						ptr.BlockPointer, NewFileBlock(), TransientEntry)
				}
			} else {
				dirEntries := dirEntriesBySizeAsc{dirEntryMapToDirEntries(b.Children)}
				sort.Sort(dirEntries)
				for i, entry := range dirEntries.dirEntries {
					// Prioritize small files
					priority := dirEntryPrefetchPriority - i
					var block Block
					switch entry.Type {
					case Dir:
						block = NewDirBlock()
					case File:
						block = NewFileBlock()
					case Exec:
						block = NewFileBlock()
					default:
						continue
					}
					p.request(priority, kmd, entry.BlockPointer,
						block, TransientEntry)
				}
			}
		}
	default:
	}
	return nil
}

func (p *blockPrefetcher) request(priority int, kmd KeyMetadata, ptr BlockPointer, block Block, lifetime BlockCacheLifetime) {
	// TODO: track these requests and do something intelligent with
	// cancellation
	ctx := context.Background()
	// Returns a buffered channel, so we don't need to read from it.
	_ = p.retriever.Request(ctx, priority, kmd, ptr, block, lifetime)
}<commit_msg>prefetcher: refactor to make the various prefetch scenarios testable<commit_after>package libkbfs

import (
	"sort"

	"golang.org/x/net/context"
)

const (
	defaultIndirectPointerPrefetchCount int = 20
	fileIndirectBlockPrefetchPriority   int = -100
	dirEntryPrefetchPriority            int = -200
)

type dirEntries []DirEntry
type dirEntriesBySizeAsc struct{ dirEntries }
type dirEntriesBySizeDesc struct{ dirEntries }

func (d dirEntries) Len() int                     { return len(d) }
func (d dirEntries) Swap(i, j int)                { d[i], d[j] = d[j], d[i] }
func (d dirEntriesBySizeAsc) Less(i, j int) bool  { return d.dirEntries[i].Size < d.dirEntries[j].Size }
func (d dirEntriesBySizeDesc) Less(i, j int) bool { return d.dirEntries[i].Size > d.dirEntries[j].Size }

func dirEntryMapToDirEntries(entryMap map[string]DirEntry) dirEntries {
	dirEntries := make(dirEntries, 0, len(entryMap))
	for _, entry := range entryMap {
		dirEntries = append(dirEntries, entry)
	}
	return dirEntries
}

type prefetcher interface {
	HandleBlock(b Block, kmd KeyMetadata, priority int)
}

var _ prefetcher = (*blockPrefetcher)(nil)

type blockPrefetcher struct {
	retriever blockRetriever
}

func newPrefetcher(retriever blockRetriever) prefetcher {
	return &blockPrefetcher{
		retriever: retriever,
	}
}

func (p *blockPrefetcher) prefetchIndirectFileBlock(b *FileBlock, kmd KeyMetadata, priority int) {
	// prefetch the first <n> indirect block pointers.
	// TODO: do something smart with subsequent blocks.
	numIPtrs := len(b.IPtrs)
	if numIPtrs > defaultIndirectPointerPrefetchCount {
		numIPtrs = defaultIndirectPointerPrefetchCount
	}
	for _, ptr := range b.IPtrs[:numIPtrs] {
		p.request(fileIndirectBlockPrefetchPriority, kmd,
			ptr.BlockPointer, NewFileBlock(), TransientEntry)
	}
}

func (p *blockPrefetcher) prefetchIndirectDirBlock(b *DirBlock, kmd KeyMetadata, priority int) {
	numIPtrs := len(b.IPtrs)
	if numIPtrs > defaultIndirectPointerPrefetchCount {
		numIPtrs = defaultIndirectPointerPrefetchCount
	}
	for _, ptr := range b.IPtrs[:numIPtrs] {
		p.request(fileIndirectBlockPrefetchPriority, kmd,
			ptr.BlockPointer, NewFileBlock(), TransientEntry)
	}
}

func (p *blockPrefetcher) prefetchDirectDirBlock(b *DirBlock, kmd KeyMetadata, priority int) {
	dirEntries := dirEntriesBySizeAsc{dirEntryMapToDirEntries(b.Children)}
	sort.Sort(dirEntries)
	for i, entry := range dirEntries.dirEntries {
		// Prioritize small files
		priority := dirEntryPrefetchPriority - i
		var block Block
		switch entry.Type {
		case Dir:
			block = NewDirBlock()
		case File:
			block = NewFileBlock()
		case Exec:
			block = NewFileBlock()
		default:
			continue
		}
		p.request(priority, kmd, entry.BlockPointer,
			block, TransientEntry)
	}
}

func (p *blockPrefetcher) HandleBlock(b Block, kmd KeyMetadata, priority int) {
	switch b := b.(type) {
	case *FileBlock:
		if b.IsInd && priority >= defaultOnDemandRequestPriority {
			p.prefetchIndirectFileBlock(b, kmd, priority)
		}
	case *DirBlock:
		// If this is an on-demand request:
		// - If the block is indirect, prefetch the first <n> indirect block
		// pointers.
		// - If the block is direct (has Children), prefetch all DirEntry root
		// blocks.
		if priority >= defaultOnDemandRequestPriority {
			if b.IsInd {
				p.prefetchIndirectDirBlock(b, kmd, priority)
			} else {
				p.prefetchDirectDirBlock(b, kmd, priority)
			}
		}
	default:
	}
}

func (p *blockPrefetcher) request(priority int, kmd KeyMetadata, ptr BlockPointer, block Block, lifetime BlockCacheLifetime) {
	// TODO: track these requests and do something intelligent with
	// cancellation
	ctx := context.Background()
	// Returns a buffered channel, so we don't need to read from it.
	_ = p.retriever.Request(ctx, priority, kmd, ptr, block, lifetime)
}
{'fraction_non_alphanumeric': 0.09936428966279713, 'fraction_numerical': 0.0024875621890547263, 'mean_word_length': 3.9738831615120276, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6097377', 'n_tokens_mistral': 2585, 'n_tokens_neox': 2345, 'n_words': 638}