fork of https://github.com/sourcegraph/zoekt
1// Copyright 2016 Google Inc. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package main
16
17import (
18 "flag"
19 "fmt"
20 "log"
21 "os"
22 "path/filepath"
23 "runtime/pprof"
24 "strings"
25
26 "github.com/sourcegraph/zoekt"
27 "github.com/sourcegraph/zoekt/build"
28 "github.com/sourcegraph/zoekt/cmd"
29 "go.uber.org/automaxprocs/maxprocs"
30)
31
32type fileInfo struct {
33 name string
34 size int64
35}
36
37type fileAggregator struct {
38 ignoreDirs map[string]struct{}
39 sizeMax int64
40 sink chan fileInfo
41}
42
43func (a *fileAggregator) add(path string, info os.FileInfo, err error) error {
44 if err != nil {
45 return err
46 }
47
48 if info.IsDir() {
49 base := filepath.Base(path)
50 if _, ok := a.ignoreDirs[base]; ok {
51 return filepath.SkipDir
52 }
53 }
54
55 if info.Mode().IsRegular() {
56 a.sink <- fileInfo{path, info.Size()}
57 }
58 return nil
59}
60
61func main() {
62 cpuProfile := flag.String("cpu_profile", "", "write cpu profile to file")
63 ignoreDirs := flag.String("ignore_dirs", ".git,.hg,.svn", "comma separated list of directories to ignore.")
64 flag.Parse()
65
66 // Tune GOMAXPROCS to match Linux container CPU quota.
67 _, _ = maxprocs.Set()
68
69 opts := cmd.OptionsFromFlags()
70 if *cpuProfile != "" {
71 f, err := os.Create(*cpuProfile)
72 if err != nil {
73 log.Fatal(err)
74 }
75 if err := pprof.StartCPUProfile(f); err != nil {
76 log.Fatal(err)
77 }
78 defer pprof.StopCPUProfile()
79 }
80
81 ignoreDirMap := map[string]struct{}{}
82 if *ignoreDirs != "" {
83 dirs := strings.Split(*ignoreDirs, ",")
84 for _, d := range dirs {
85 d = strings.TrimSpace(d)
86 if d != "" {
87 ignoreDirMap[d] = struct{}{}
88 }
89 }
90 }
91 for _, arg := range flag.Args() {
92 opts.RepositoryDescription.Source = arg
93 if err := indexArg(arg, *opts, ignoreDirMap); err != nil {
94 log.Fatal(err)
95 }
96 }
97}
98
99func indexArg(arg string, opts build.Options, ignore map[string]struct{}) error {
100 dir, err := filepath.Abs(filepath.Clean(arg))
101 if err != nil {
102 return err
103 }
104
105 opts.RepositoryDescription.Name = filepath.Base(dir)
106 builder, err := build.NewBuilder(opts)
107 if err != nil {
108 return err
109 }
110 // we don't need to check error, since we either already have an error, or
111 // we returning the first call to builder.Finish.
112 defer builder.Finish() // nolint:errcheck
113
114 comm := make(chan fileInfo, 100)
115 agg := fileAggregator{
116 ignoreDirs: ignore,
117 sink: comm,
118 sizeMax: int64(opts.SizeMax),
119 }
120
121 go func() {
122 if err := filepath.Walk(dir, agg.add); err != nil {
123 log.Fatal(err)
124 }
125 close(comm)
126 }()
127
128 for f := range comm {
129 displayName := strings.TrimPrefix(f.name, dir+"/")
130 if f.size > int64(opts.SizeMax) && !opts.IgnoreSizeMax(displayName) {
131 if err := builder.Add(zoekt.Document{
132 Name: displayName,
133 SkipReason: fmt.Sprintf("document size %d larger than limit %d", f.size, opts.SizeMax),
134 }); err != nil {
135 return err
136 }
137 continue
138 }
139 content, err := os.ReadFile(f.name)
140 if err != nil {
141 return err
142 }
143
144 if err := builder.AddFile(displayName, content); err != nil {
145 return err
146 }
147 }
148
149 return builder.Finish()
150}