-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
186 lines (169 loc) · 4.26 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
package main
import (
"flag"
"fmt"
"os"
"runtime"
"strings"
"sync"
)
var IGNORE_PATHS = [7]string{
".git",
"build",
"dist",
"env",
"node_modules",
"target",
"venv",
}
var (
version = "development"
commit = "unknown"
date = "2024-01-01"
builtBy = "lwileczek"
)
func printBuildInfo() {
fmt.Printf("Version: %s\nCommit: %s\nBuild Date: %s | by: %s\n", version, commit, date, builtBy)
}
func main() {
//The number of workers. If there are more workers the system can read from
//the work queue more often and a larger queue is not required.
workers := flag.Int("w", -1, "Number of workers")
//If the queue overflows we'll use a slice to store work which might slow the system
queueSize := flag.Int("q", 512, "The max work queue size")
maxResults := flag.Int("c", -1, "The maximum number of results to find")
dir := flag.String("d", ".", "The starting directory to check for files")
pattern := flag.String("p", "", "A pattern to check for within the file names")
v := flag.Bool("v", false, "print the version and build information")
flag.Parse()
if *v {
printBuildInfo()
return
}
if *pattern == "" {
fmt.Println("No pattern provided")
return
}
w := *workers
if *workers <= 0 {
w = runtime.NumCPU() + 2
}
//Only for OSX/Linux, sorry windows
//Remove any trailing slashes in the path
if (*dir)[len(*dir)-1:] == "/" {
*dir = string((*dir)[0 : len(*dir)-1])
}
printCh := make(chan string, w)
//The system will reach deadlock if the work queue reaches capacity
workQ := make(chan string, *queueSize)
//To avoid deadlock, send tasks here which will have a non-blocky retry
//func to add tasks back to workQ
failover := make(chan string)
dirCount := make(chan int)
//Track how many dirs are open and close the work queue when we hit zero
go dirChecker(dirCount, workQ)
//Not closing as goroutines will continue to try and write if we exit early
//with the -c flag but these should be fine and killed when the program exits
//defer close(dirCount)
//defer close(failover)
go handleFailover(workQ, failover)
go createWorkerPool(pattern, workQ, failover, printCh, dirCount, w)
//Send first work request
workQ <- *dir
//Print all results
showResults(printCh, maxResults)
}
func showResults(ch chan string, limit *int) {
if *limit > 0 {
n := 0
for item := range ch {
fmt.Println(item)
n++
if n >= *limit {
return
}
}
} else {
for item := range ch {
fmt.Println(item)
}
}
}
func dirChecker(in chan int, work chan string) {
n := 1
for i := range in {
n += i
if n <= 0 {
close(work)
return
}
}
}
func createWorkerPool(p *string, in chan string, failover chan string, results chan string, cnt chan int, w int) {
var wg sync.WaitGroup
for i := 0; i < w; i++ {
wg.Add(1)
go func() {
defer wg.Done()
search(p, in, failover, results, cnt)
}()
}
wg.Wait()
close(results)
}
func search(pattern *string, in chan string, failover chan string, out chan string, cnt chan int) {
for path := range in {
items, err := os.ReadDir(path)
if err != nil {
fmt.Println("Error reading the directory", path)
fmt.Println(err)
cnt <- -1
}
ItemSearch:
for _, item := range items {
if item.IsDir() {
//Don't dive into directories I don't care about
for _, p := range IGNORE_PATHS {
if p == item.Name() {
continue ItemSearch
}
}
subPath := fmt.Sprintf("%s/%s", path, item.Name())
cnt <- 1
select {
case in <- subPath:
case failover <- subPath:
}
}
//Always check if the name of the thing matches pattern, including directory names
if strings.Index(item.Name(), *pattern) >= 0 {
//subPath is repeated but no point in creating an allocation if not required
subPath := fmt.Sprintf("%s/%s", path, item.Name())
out <- subPath
}
}
//We finished reading everything in the dir, tell the accounted we finished
cnt <- -1
}
}
func handleFailover(work, fail chan string) {
var q []string
for {
task := <-fail
q = append(q, task)
//TODO: Add verbose logging here so users can check if the failover was used
for {
select {
case work <- q[0]:
q = q[1:]
case task := <-fail:
q = append(q, task)
default:
}
//I don't know if we'll get an issue with `work <- q[0]` unless we have this
if len(q) == 0 {
break
}
}
}
}