From cfe7ca7cea0eb656dda4df1ef76c610edd0d7b6e Mon Sep 17 00:00:00 2001 From: Bryan Boreham Date: Tue, 4 Jun 2019 17:07:58 +0000 Subject: [PATCH] Read a file of org-hour records --- extras/scanner/main.go | 59 ++++++++++++++++++++++++++++++------------ 1 file changed, 43 insertions(+), 16 deletions(-) diff --git a/extras/scanner/main.go b/extras/scanner/main.go index bf2eee79f..c535d8939 100644 --- a/extras/scanner/main.go +++ b/extras/scanner/main.go @@ -1,6 +1,7 @@ package main import ( + "bufio" "context" "flag" "fmt" @@ -83,7 +84,8 @@ func main() { queryRateLimit float64 writeRateLimit float64 - orgsFile string + orgsFile string + recordsFile string scanner scanner loglevel string @@ -102,6 +104,7 @@ func main() { flag.IntVar(&scanner.segments, "segments", 1, "Number of segments to read in parallel") flag.StringVar(&scanner.address, "address", ":6060", "Address to listen on, for profiling, etc.") flag.StringVar(&orgsFile, "delete-orgs-file", "", "File containing IDs of orgs to delete") + flag.StringVar(&recordsFile, "delete-records-file", "", "File containing IDs of orgs to delete") flag.StringVar(&loglevel, "log-level", "info", "Debug level: debug, info, warning, error") flag.BoolVar(&justBigScan, "big-scan", false, "If true, just scan the whole index and print summaries") @@ -141,12 +144,7 @@ func main() { return } - orgs := []string{} - if orgsFile != "" { - content, err := ioutil.ReadFile(orgsFile) - checkFatal(err) - orgs = strings.Fields(string(content)) - } + orgs := setupOrgs(orgsFile) if scanner.stopHour == 0 { scanner.stopHour = int(time.Now().Unix() / int64(time.Hour/time.Second)) @@ -158,21 +156,50 @@ func main() { totals := newSummary() - var orgWait sync.WaitGroup - orgWait.Add(len(orgs)) - - for _, org := range orgs { - go func(org string) { - scanner.processOrg(context.Background(), org) - orgWait.Done() - }(org) + if recordsFile != "" { + f, err := os.Open(recordsFile) + checkFatal(err) + defer f.Close() + records := bufio.NewScanner(f) + for records.Scan() { + scanner.HandleRecord(context.Background(), orgs, records.Text()) + } + checkFatal(records.Err()) } - orgWait.Wait() fmt.Printf("\n") totals.print() } +func setupOrgs(deleteOrgsFile string) (deleteOrgs map[int]struct{}) { + deleteOrgs = map[int]struct{}{} + if deleteOrgsFile != "" { + content, err := ioutil.ReadFile(deleteOrgsFile) + checkFatal(err) + for _, arg := range strings.Fields(string(content)) { + org, err := strconv.Atoi(arg) + checkFatal(err) + deleteOrgs[org] = struct{}{} + } + } + return +} + +func (sc *scanner) HandleRecord(ctx context.Context, deleteOrgs map[int]struct{}, record string) { + fields := strings.Split(record, "-") + org, err := strconv.Atoi(fields[0]) + checkFatal(err) + if _, found := deleteOrgs[org]; !found { + return + } + hour, err := strconv.Atoi(fields[1]) + checkFatal(err) + if hour < sc.startHour || hour > sc.stopHour { + return + } + sc.deleteOneOrgHour(ctx, fields[0], hour) +} + func (sc *scanner) processOrg(ctx context.Context, org string) { deleted := 0 for hour := sc.startHour; hour <= sc.stopHour; hour++ {