fork of https://github.com/sourcegraph/zoekt
1package main
2
3import (
4 "bufio"
5 "fmt"
6 "log"
7 "os"
8 "path/filepath"
9 "strings"
10
11 "github.com/sourcegraph/zoekt"
12)
13
14// merge merges the input shards into a compound shard in dstDir. It returns the
15// full path to the compound shard. The input shards are removed on success.
16func merge(dstDir string, names []string) (string, error) {
17 var files []zoekt.IndexFile
18 for _, fn := range names {
19 f, err := os.Open(fn)
20 if err != nil {
21 return "", nil
22 }
23 defer f.Close()
24
25 indexFile, err := zoekt.NewIndexFile(f)
26 if err != nil {
27 return "", err
28 }
29 defer indexFile.Close()
30
31 files = append(files, indexFile)
32 }
33
34 tmpName, dstName, err := zoekt.Merge(dstDir, files...)
35 if err != nil {
36 return "", err
37 }
38
39 // Delete input shards.
40 for _, name := range names {
41 paths, err := zoekt.IndexFilePaths(name)
42 if err != nil {
43 return "", fmt.Errorf("zoekt-merge-index: %w", err)
44 }
45 for _, p := range paths {
46 if err := os.Remove(p); err != nil {
47 return "", fmt.Errorf("zoekt-merge-index: failed to remove simple shard: %w", err)
48 }
49 }
50 }
51
52 // We only rename the compound shard if all simple shards could be deleted in the
53 // previous step. This guarantees we won't have duplicate indexes.
54 if err := os.Rename(tmpName, dstName); err != nil {
55 return "", fmt.Errorf("zoekt-merge-index: failed to rename compound shard: %w", err)
56 }
57
58 return dstName, nil
59}
60
61func mergeCmd(paths []string) (string, error) {
62 if paths[0] == "-" {
63 paths = []string{}
64 scanner := bufio.NewScanner(os.Stdin)
65 for scanner.Scan() {
66 paths = append(paths, strings.TrimSpace(scanner.Text()))
67 }
68 if err := scanner.Err(); err != nil {
69 return "", err
70 }
71 log.Printf("merging %d paths from stdin", len(paths))
72 }
73
74 return merge(filepath.Dir(paths[0]), paths)
75}
76
77// explode splits the input shard into individual shards and places them in dstDir.
78// Temporary files created in the process are removed on a best effort basis.
79func explode(dstDir string, inputShard string) error {
80 f, err := os.Open(inputShard)
81 if err != nil {
82 return err
83 }
84 defer f.Close()
85
86 indexFile, err := zoekt.NewIndexFile(f)
87 if err != nil {
88 return err
89 }
90 defer indexFile.Close()
91
92 exploded, err := zoekt.Explode(dstDir, indexFile)
93 defer func() {
94 // best effort removal of tmp files. If os.Remove fails, indexserver will delete
95 // the leftover tmp files during the next cleanup.
96 for tmpFn := range exploded {
97 os.Remove(tmpFn)
98 }
99 }()
100 if err != nil {
101 return fmt.Errorf("zoekt.Explode: %w", err)
102 }
103
104 // remove the input shard first to avoid duplicate indexes. In the worst case,
105 // the process is interrupted just after we delete the compound shard, in which
106 // case we have to reindex the lost repos.
107 paths, err := zoekt.IndexFilePaths(inputShard)
108 if err != nil {
109 return err
110 }
111 for _, path := range paths {
112 err = os.Remove(path)
113 if err != nil {
114 return err
115 }
116 }
117
118 // best effort rename shards.
119 for tmpFn, dstFn := range exploded {
120 if err := os.Rename(tmpFn, dstFn); err != nil {
121 log.Printf("explode: rename failed: %s", err)
122 }
123 }
124
125 return nil
126}
127
128func explodeCmd(path string) error {
129 return explode(filepath.Dir(path), path)
130}
131
132func main() {
133 switch subCommand := os.Args[1]; subCommand {
134 case "merge":
135 compoundShardPath, err := mergeCmd(os.Args[2:])
136 if err != nil {
137 log.Fatal(err)
138 }
139 fmt.Println(compoundShardPath)
140 case "explode":
141 if err := explodeCmd(os.Args[2]); err != nil {
142 log.Fatal(err)
143 }
144 default:
145 log.Fatalf("unknown subcommand %s", subCommand)
146 }
147}