Charlie GunyonandCharlie Gunyon
311494e86e
* Add stub source and elastic API funcs * Spawn workers and ship chunks * Now successfully detects a credential - Added tests - Added some documentation comments - Threaded the passed context through to all the API requests * Linting fixes * Add integration tests and resolve some bugs they uncovered * Logstash -> Elasticsearch * Add support for --index-pattern * Add support for --query-json * Use structs instead of string building to construct a search body * Support --since-timestamp * Implement additional authentication methods * Fix some small bugs * Refactoring to support --best-effort-scan * Finish implementation of --best-effort-scan * Implement scan catch-up * Finish connecting support for nodes CLI arg * Add some integration tests around the catchup mechanism * go mod tidy * Fix some linting issues * Remove some debugging Prints * Move off of _doc * Remove informational Printf and add informational logging * Remove debugging logging * Copy the index from the outer loop as well * Don't burn up the ES API with rapid requests if there's no work to do in subsequent scans * No need to export UnitOfWork.AddSearch * Use a better name for the range query variable when building the timestamp range clause in searches * Replace some unlocking defers with explicit unlocks to make the synchronized part of the code clearer * found -> ok * Remove superfluous buildElasticClient method --------- Co-authored-by: Charlie Gunyon <[email protected]>
102 lines
2.4 KiB
Go
102 lines
2.4 KiB
Go
package elasticsearch
|
|
|
|
import "fmt"
|
|
|
|
type DocumentSearch struct {
|
|
index *Index
|
|
offset int
|
|
documentCount int
|
|
skipCount int
|
|
filterParams *FilterParams
|
|
}
|
|
|
|
type UnitOfWork struct {
|
|
maxDocumentCount int
|
|
documentCount int
|
|
documentSearches []DocumentSearch
|
|
}
|
|
|
|
func NewUnitOfWork(maxDocumentCount int) UnitOfWork {
|
|
uow := UnitOfWork{maxDocumentCount: maxDocumentCount}
|
|
uow.documentSearches = []DocumentSearch{}
|
|
|
|
return uow
|
|
}
|
|
|
|
func (ds *DocumentSearch) String() string {
|
|
if ds.offset > 0 {
|
|
return fmt.Sprintf("%s [%d:]", ds.index.name, ds.offset)
|
|
} else {
|
|
return ds.index.name
|
|
}
|
|
}
|
|
|
|
func (uow *UnitOfWork) addSearch(
|
|
index *Index,
|
|
filterParams *FilterParams,
|
|
offset int,
|
|
scanCoverageRate float64,
|
|
) int {
|
|
indexDocCount := index.documentCount - offset
|
|
addedDocumentCount := min(uow.maxDocumentCount-uow.documentCount, indexDocCount)
|
|
|
|
if addedDocumentCount > 0 {
|
|
uow.documentSearches = append(uow.documentSearches, DocumentSearch{
|
|
index: index,
|
|
offset: offset,
|
|
documentCount: addedDocumentCount,
|
|
skipCount: int(float64(addedDocumentCount) * (1.0 - scanCoverageRate)),
|
|
filterParams: filterParams,
|
|
})
|
|
|
|
uow.documentCount += addedDocumentCount
|
|
}
|
|
|
|
return addedDocumentCount
|
|
}
|
|
|
|
func distributeDocumentScans(
|
|
indices *Indices,
|
|
maxUnits int,
|
|
scanCoverageRate float64,
|
|
) []UnitOfWork {
|
|
totalDocumentCount := 0
|
|
|
|
for _, i := range indices.indices {
|
|
totalDocumentCount += i.documentCount
|
|
}
|
|
|
|
unitsOfWork := make([]UnitOfWork, maxUnits)
|
|
documentsAssigned := 0
|
|
for i := 0; i < maxUnits; i++ {
|
|
documentCount := totalDocumentCount / maxUnits
|
|
|
|
// The total number of documents to process might not be perfectly
|
|
// divisible by the number of workers, so make sure any remaining documents
|
|
// get processed by assigning them to the last worker
|
|
if i == maxUnits-1 {
|
|
documentCount = totalDocumentCount - documentsAssigned
|
|
}
|
|
|
|
unitsOfWork[i] = NewUnitOfWork(documentCount)
|
|
|
|
documentsAssigned += documentCount
|
|
}
|
|
|
|
unitOfWorkIndex := 0
|
|
for _, i := range indices.indices {
|
|
uow := &unitsOfWork[unitOfWorkIndex]
|
|
offset := uow.addSearch(i, indices.filterParams, 0, scanCoverageRate)
|
|
|
|
// If we've yet to distribute all the documents in the index, go into the
|
|
// next unit of work, and the next, and the next....
|
|
for offset < i.documentCount {
|
|
unitOfWorkIndex++
|
|
uow := &unitsOfWork[unitOfWorkIndex]
|
|
offset += uow.addSearch(i, indices.filterParams, offset, scanCoverageRate)
|
|
}
|
|
}
|
|
|
|
return unitsOfWork
|
|
}
|