fileslist.go 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. // Copyright 2017 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. // fileslist.py replacement written in GO, which utilizes multi-cores.
  15. package main
  16. import (
  17. "crypto/sha256"
  18. "encoding/json"
  19. "flag"
  20. "fmt"
  21. "io"
  22. "os"
  23. "path/filepath"
  24. "runtime"
  25. "sort"
  26. "strings"
  27. "sync"
  28. )
  29. const (
  30. MAX_DEFAULT_PARA = 24
  31. )
  32. func defaultPara() int {
  33. ret := runtime.NumCPU()
  34. if ret > MAX_DEFAULT_PARA {
  35. return MAX_DEFAULT_PARA
  36. }
  37. return ret
  38. }
  39. var (
  40. para = flag.Int("para", defaultPara(), "Number of goroutines")
  41. )
  42. // Represents each file.
  43. type Node struct {
  44. SHA256 string
  45. Name string // device side path.
  46. Size int64
  47. path string // host side path.
  48. stat os.FileInfo
  49. }
  50. func newNode(hostPath string, devicePath string, stat os.FileInfo) Node {
  51. return Node{Name: devicePath, path: hostPath, stat: stat}
  52. }
  53. // Scan a Node and returns true if it should be added to the result.
  54. func (n *Node) scan() bool {
  55. n.Size = n.stat.Size()
  56. // Calculate SHA256.
  57. h := sha256.New()
  58. if n.stat.Mode()&os.ModeSymlink == 0 {
  59. f, err := os.Open(n.path)
  60. if err != nil {
  61. panic(err)
  62. }
  63. defer f.Close()
  64. if _, err := io.Copy(h, f); err != nil {
  65. panic(err)
  66. }
  67. } else {
  68. // Hash the content of symlink, not the file it points to.
  69. s, err := os.Readlink(n.path)
  70. if err != nil {
  71. panic(err)
  72. }
  73. if _, err := io.WriteString(h, s); err != nil {
  74. panic(err)
  75. }
  76. }
  77. n.SHA256 = fmt.Sprintf("%x", h.Sum(nil))
  78. return true
  79. }
  80. func main() {
  81. flag.Parse()
  82. allOutput := make([]Node, 0, 1024) // Store all outputs.
  83. mutex := &sync.Mutex{} // Guard allOutput
  84. ch := make(chan Node) // Pass nodes to goroutines.
  85. var wg sync.WaitGroup // To wait for all goroutines.
  86. wg.Add(*para)
  87. // Scan files in multiple goroutines.
  88. for i := 0; i < *para; i++ {
  89. go func() {
  90. defer wg.Done()
  91. output := make([]Node, 0, 1024) // Local output list.
  92. for node := range ch {
  93. if node.scan() {
  94. output = append(output, node)
  95. }
  96. }
  97. // Add to the global output list.
  98. mutex.Lock()
  99. allOutput = append(allOutput, output...)
  100. mutex.Unlock()
  101. }()
  102. }
  103. // Walk the directories and find files to scan.
  104. for _, dir := range flag.Args() {
  105. absDir, err := filepath.Abs(dir)
  106. if err != nil {
  107. panic(err)
  108. }
  109. deviceRoot := filepath.Clean(absDir + "/..")
  110. err = filepath.Walk(dir, func(path string, stat os.FileInfo, err error) error {
  111. if err != nil {
  112. panic(err)
  113. }
  114. if stat.IsDir() {
  115. return nil
  116. }
  117. absPath, err := filepath.Abs(path)
  118. if err != nil {
  119. panic(err)
  120. }
  121. devicePath, err := filepath.Rel(deviceRoot, absPath)
  122. if err != nil {
  123. panic(err)
  124. }
  125. devicePath = "/" + devicePath
  126. ch <- newNode(absPath, devicePath, stat)
  127. return nil
  128. })
  129. if err != nil {
  130. panic(err)
  131. }
  132. }
  133. // Wait until all the goroutines finish.
  134. close(ch)
  135. wg.Wait()
  136. // Sort the entries and dump as json.
  137. sort.Slice(allOutput, func(i, j int) bool {
  138. if allOutput[i].Size > allOutput[j].Size {
  139. return true
  140. }
  141. if allOutput[i].Size == allOutput[j].Size && strings.Compare(allOutput[i].Name, allOutput[j].Name) > 0 {
  142. return true
  143. }
  144. return false
  145. })
  146. j, err := json.MarshalIndent(allOutput, "", " ")
  147. if err != nil {
  148. panic(nil)
  149. }
  150. fmt.Printf("%s\n", j)
  151. }