-
Notifications
You must be signed in to change notification settings - Fork 1
/
pcopy.go
138 lines (115 loc) · 2.79 KB
/
pcopy.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
package main
import (
"crypto/sha1"
"flag"
"fmt"
"io"
"log"
"os"
"runtime"
"sync"
)
var maxWorkers int = 30
// pcopy copies files from the top-level of a source directory to a
// destination directory using up to 30 concurrent goroutines. The
// default number of goroutines is 10. The destination directory must
// exist.
func main() {
var numWorkers = flag.Int("w", 10, "number of workers")
flag.Parse()
if len(flag.Args()) < 2 {
log.Fatalf("usage: %s [ -w ] srcDir destDir", os.Args[0])
}
log.Println("GOMAXPROCS:", runtime.GOMAXPROCS(0))
log.Println("NumCPU:", runtime.NumCPU())
srcDir := flag.Args()[0]
destDir := flag.Args()[1]
if *numWorkers > 30 {
log.Println(*numWorkers, "specifed; max is", maxWorkers)
*numWorkers = maxWorkers
}
log.Println("number of workers:", *numWorkers)
todo := make(chan string)
results := make(chan string, *numWorkers)
var wg sync.WaitGroup
if !isThere(destDir) {
log.Fatal("Can't find directory ", destDir)
}
files := filesIn(srcDir)
// queue up files to work on
go func(fs []string) {
for _, f := range fs {
todo <- f
}
close(todo)
}(files)
// start the workers
for i := 0; i < *numWorkers; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
for f := range todo {
srcName := fmt.Sprintf("%s/%s", srcDir, f)
destName := fmt.Sprintf("%s/%s", destDir, f)
numBytes, chksum := copyAndSha1Sum(srcName, destName)
results <- fmt.Sprintf("%s %d %s", f, numBytes, chksum)
}
}(i + 1)
}
// when workers are done, close the results channel
go func() {
wg.Wait()
close(results)
}()
// get the results
for line := range results {
fmt.Println(line)
}
}
// Check for existence of dir
func isThere(dir string) bool {
if _, err := os.Open(dir); err != nil {
return false
}
return true
}
// Return a list of files in the top level of dir
func filesIn(dir string) []string {
var fs []string
d, _ := os.Open(dir)
names, err := d.Readdirnames(0)
if err != nil {
log.Fatal("filesIn: ", err)
}
for _, name := range names {
finfo, err := os.Lstat(dir + "/" + name)
if err != nil {
log.Fatal("filesIn: ", err)
}
if finfo.Mode().IsRegular() {
fs = append(fs, finfo.Name())
}
}
return fs
}
// Copy file from src to dest and compute SHA1 sum of bytes copied
// Returns the number of bytes written and the checksum
func copyAndSha1Sum(srcName string, destName string) (int64, string) {
src, err := os.Open(srcName)
if err != nil {
log.Fatal("copyAndSha1Sum: ", err)
}
defer src.Close()
dest, err := os.Create(destName)
if err != nil {
log.Fatal("copyAndSha1Sum: ", err)
}
defer dest.Close()
h := sha1.New()
tee := io.MultiWriter(dest, h)
bytesWritten, err := io.Copy(tee, src)
if err != nil {
log.Fatal("copyAndSha1Sum: ", err)
}
return bytesWritten, fmt.Sprintf("%x", h.Sum(nil))
}