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