-/*
-glocate -- ZFS-diff-friendly locate-like utility
-Copyright (C) 2022 Sergey Matveev <stargrave@stargrave.org>
-
-This program is free software: you can redistribute it and/or modify
-it under the terms of the GNU General Public License as published by
-the Free Software Foundation, version 3 of the License.
-
-This program is distributed in the hope that it will be useful,
-but WITHOUT ANY WARRANTY; without even the implied warranty of
-MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-GNU General Public License for more details.
-
-You should have received a copy of the GNU General Public License
-along with this program. If not, see <http://www.gnu.org/licenses/>.
-*/
-
package main
import (
"bufio"
- "encoding/gob"
"flag"
- "fmt"
- "io"
- "io/fs"
"log"
"os"
- "path"
- "sort"
- "strconv"
"strings"
"syscall"
- "time"
-
- "github.com/dustin/go-humanize"
- "github.com/klauspost/compress/zstd"
)
-type File struct {
- Name string
- Size uint64
- Mtime int64
- Files []File
+type Ent struct {
+ name []string
+ mtime int64
+ size int64
}
-type ByName []File
-
-func (a ByName) Len() int {
- return len(a)
-}
-
-func (a ByName) Swap(i, j int) {
- a[i], a[j] = a[j], a[i]
-}
-
-func (a ByName) Less(i, j int) bool {
- return a[i].Name < a[j].Name
-}
-
-func (file *File) IsDir() bool {
- return file.Name[len(file.Name)-1] == '/'
-}
-
-func walk(root string) ([]File, uint64, error) {
- fd, err := os.Open(root)
- if err != nil {
- return nil, 0, err
- }
- var files []File
- var size uint64
- var info fs.FileInfo
- for {
- ents, err := fd.ReadDir(1 << 10)
- if err != nil {
- if err == io.EOF {
- break
- }
- fd.Close()
- return nil, 0, err
- }
- for _, ent := range ents {
- file := File{Name: ent.Name()}
- fullPath := path.Join(root, file.Name)
- if ent.IsDir() {
- file.Name += "/"
- }
- info, err = ent.Info()
- if err != nil {
- log.Println("can not stat:", fullPath, ":", err)
- files = append(files, file)
- continue
- }
- file.Mtime = info.ModTime().Unix()
- if ent.IsDir() {
- file.Files, file.Size, err = walk(fullPath)
- if err != nil {
- log.Println("can not walk:", fullPath, ":", err)
- files = append(files, file)
- continue
- }
- } else if info.Mode().IsRegular() {
- file.Size = uint64(info.Size())
- }
- files = append(files, file)
- size += file.Size
- }
- }
- fd.Close()
- sort.Sort(ByName(files))
- return files, size, nil
-}
-
-func usage() {
- log.Println("usage")
- os.Exit(1)
-}
-
-func load(dbPath string) *File {
- fd, err := os.Open(dbPath)
- if err != nil {
- log.Fatalln(err)
- }
- defer fd.Close()
- comp, err := zstd.NewReader(fd)
- if err != nil {
- log.Fatalln(err)
- }
- dec := gob.NewDecoder(comp)
- var file File
- err = dec.Decode(&file)
- if err != nil {
- log.Fatalln(err)
- }
- comp.Close()
- return &file
+func (ent *Ent) IsDir() bool {
+ return IsDir(ent.name[len(ent.name)-1])
}
-func (db *File) dump(dbPath string) error {
- tmp, err := os.CreateTemp(path.Dir(dbPath), "glocate")
- if err != nil {
- return err
- }
- defer os.Remove(tmp.Name())
- comp, err := zstd.NewWriter(
- tmp, zstd.WithEncoderLevel(zstd.SpeedBestCompression),
- )
- if err != nil {
- return err
- }
- enc := gob.NewEncoder(comp)
- err = enc.Encode(db)
- if err != nil {
- return err
- }
- err = comp.Close()
- if err != nil {
- return err
- }
- err = tmp.Close()
- if err != nil {
- return err
- }
+func dbCommit(dbPath string, tmp *os.File) {
umask := syscall.Umask(0)
syscall.Umask(umask)
- err = os.Chmod(tmp.Name(), os.FileMode(0666&^umask))
- if err != nil {
- return err
- }
- return os.Rename(tmp.Name(), dbPath)
-}
-
-func (file *File) listBeauty(indent string, n int, isLast, veryFirst bool) {
- if veryFirst {
- fmt.Printf("[%s]\n", humanize.IBytes(file.Size))
- } else {
- var box string
- if isLast {
- box = "â"
- } else {
- box = "â"
- }
- name := file.Name
- fmt.Printf("%s%s %s\tâ%d [%s] %s\n",
- indent, box, name, n, humanize.IBytes(file.Size),
- time.Unix(file.Mtime, 0).Format("2006-01-02"),
- )
- if isLast {
- indent += " "
- } else {
- indent += "â "
- }
- }
- for n, f := range file.Files {
- n++
- f.listBeauty(indent, n, n == len(file.Files), false)
- }
-}
-
-func (file *File) listSimple(root string, veryFirst bool) {
- name := file.Name
- fmt.Println(
- strconv.FormatUint(file.Size, 10),
- time.Unix(file.Mtime, 0).Format("2006-01-02T15:04:05"),
- root+name,
- )
- if veryFirst {
- name = ""
- }
- for _, f := range file.Files {
- f.listSimple(root+name, false)
- }
-}
-
-func (file *File) listFiles(root string, veryFirst bool) {
- name := file.Name
- if veryFirst {
- root = ""
- } else {
- fmt.Println(root + name)
- root += name
- }
- for _, f := range file.Files {
- f.listFiles(root, false)
- }
-}
-
-func (db *File) find(p string) (file *File, parents []*File, idx int, err error) {
- file = db
- var f File
-Entities:
- for _, ent := range strings.Split(p, "/") {
- for idx, f = range file.Files {
- if (ent == f.Name) || (ent+"/" == f.Name) {
- parents = append(parents, file)
- file = &f
- continue Entities
- }
- }
- err = fmt.Errorf("no entity found: %s", ent)
- return
- }
- return
-}
-
-func (db *File) remove(p string) error {
- file, parents, idx, err := db.find(p)
- if err != nil {
- return err
- }
- lastParent := parents[len(parents)-1]
- lastParent.Files = append(
- lastParent.Files[:idx],
- lastParent.Files[idx+1:]...,
- )
- for _, parent := range parents {
- parent.Size -= file.Size
- }
- return nil
-}
-
-func (db *File) add(p string) error {
- cols := strings.Split(p, "/")
- cols, name := cols[:len(cols)-1], cols[len(cols)-1]
- var parent *File
- var err error
- if len(cols) != 0 {
- parent, _, _, err = db.find(path.Join(cols...))
- if err != nil {
- return err
- }
- } else {
- parent = db
- }
- info, err := os.Stat(p)
- if err != nil {
- return err
- }
- if info.IsDir() {
- name += "/"
- }
- file := File{
- Name: name,
- Size: uint64(info.Size()),
- Mtime: info.ModTime().Unix(),
+ if err := os.Chmod(tmp.Name(), os.FileMode(0666&^umask)); err != nil {
+ log.Fatalln(err)
}
- parent.Files = append(parent.Files, file)
- sort.Sort(ByName(parent.Files))
- parent.Size += file.Size
- return nil
-}
-
-func deoctalize(s string) string {
- chars := make([]byte, 0, len(s))
- for i := 0; i < len(s); i++ {
- if s[i] == '\\' {
- b, err := strconv.ParseUint("0"+s[i+1:i+1+3], 0, 8)
- if err != nil {
- log.Fatalln(err)
- }
- chars = append(chars, byte(b))
- i += 3
- } else {
- chars = append(chars, s[i])
- }
+ if err := os.Rename(tmp.Name(), dbPath); err != nil {
+ log.Fatalln(err)
}
- return string(chars)
}
func main() {
- dbPath := flag.String("db", ".glocate.db", "Path to state file (database)")
- doIndex := flag.Bool("index", false, "Initialize database")
- doUpdate := flag.Bool("update", false, "Update database by zfs-diff's output")
- showBeauty := flag.Bool("show-beauty", false, "Show beauty human-friendly listing")
- showSimple := flag.Bool("show-simple", false, "Show simple listing")
- stripPrefix := flag.String("strip-prefix", "", "Strip prefix from zfs-diff's output")
+ dbPath := flag.String("db", "glocate.db", "Path to database")
+ doIndex := flag.Bool("index", false, "Perform indexing")
+ doUpdate := flag.Bool("update", false, "Feed zfs-diff and update the database")
+ strip := flag.String("strip", "", "Strip prefix from zfs-diff's paths")
+ showMachine := flag.Bool("machine", false, "Show machine friendly")
+ showTree := flag.Bool("tree", false, "Show human-friendly tree")
+ dryRun := flag.Bool("n", false, "Dry run, do not overwrite database")
flag.Parse()
log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds | log.Lshortfile)
if *doIndex {
- files, size, err := walk(".")
- if err != nil {
- log.Fatalln(err)
- }
- db := File{Name: "./", Size: size, Files: files}
- err = db.dump(*dbPath)
- if err != nil {
- log.Fatalln(err)
+ tmp := index()
+ tmp.Close()
+ if !*dryRun {
+ dbCommit(*dbPath, tmp)
}
return
}
- db := load(*dbPath)
if *doUpdate {
- scanner := bufio.NewScanner(os.Stdin)
- var t string
- for scanner.Scan() {
- t = scanner.Text()
- if len(t) == 0 {
- continue
- }
- cols := strings.Split(t, "\t")
- if len(cols) < 2 {
- log.Fatalln("bad zfs-diff format")
- }
- switch cols[0] {
- case "-":
- name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
- if err := db.remove(name); err != nil {
- log.Println("can not -:", name, ":", err)
- }
- case "+":
- name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
- if err := db.add(name); err != nil {
- log.Println("can not +:", name, ":", err)
- }
- case "M":
- name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
- if name == "" {
- continue
- }
- file, _, _, err := db.find(name)
- if err != nil {
- log.Println("can not M:", name, ":", err)
- continue
- }
- info, err := os.Stat(name)
- if err != nil {
- log.Println("can not M:", name, ":", err)
- continue
- }
- if info.Mode().IsRegular() {
- file.Size = uint64(info.Size())
- }
- file.Mtime = info.ModTime().Unix()
- case "R":
- if len(cols) != 3 {
- log.Fatalln("bad zfs-diff format for R")
- }
- name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
- if err := db.remove(name); err != nil {
- log.Println("can not R-:", name, ":", err)
- continue
- }
- name = deoctalize(strings.TrimPrefix(cols[2], *stripPrefix))
- if err := db.add(name); err != nil {
- log.Println("can not R+:", name, ":", err)
- }
- default:
- log.Fatalln("bad zfs-diff format")
- }
- }
- if err := scanner.Err(); err != nil {
- log.Fatalln(err)
- }
- if err := db.dump(*dbPath); err != nil {
- log.Fatalln(err)
+ tmp := updateWithDiff(*dbPath, *strip)
+ tmp.Close()
+ if !*dryRun {
+ dbCommit(*dbPath, tmp)
}
return
}
- veryFirst := true
- if len(flag.Args()) > 0 {
- root := flag.Args()[0]
- if root[:2] == "./" {
- root = root[2:]
- }
- if root[len(root)-1:] == "/" {
- root = root[:len(root)-1]
- }
- file, _, _, err := db.find(root)
- if err != nil {
- log.Fatalln(err)
- }
- db = file
- db.Name = root + "/"
- veryFirst = false
+ db, err := os.Open(*dbPath)
+ if err != nil {
+ log.Fatalln(err)
}
+ entsReader := make(chan Ent, 1<<10)
+ go reader(bufio.NewReaderSize(db, 1<<17), entsReader)
- if *showBeauty {
- db.listBeauty("", 0, false, veryFirst)
- return
+ entsPrinter := make(chan Ent, 1<<10)
+ printerJob := make(chan struct{})
+ go func() {
+ if *showMachine {
+ printerMachine(entsPrinter)
+ } else if *showTree {
+ printerTree(entsPrinter)
+ } else {
+ printerSimple(entsPrinter)
+ }
+ close(printerJob)
+ }()
+
+ var root []string
+ if len(flag.Args()) > 0 {
+ root = strings.Split("./"+flag.Arg(0), "/")
}
- if *showSimple {
- db.listSimple("", veryFirst)
- return
+
+ rootMet := false
+ for ent := range entsReader {
+ if hasPrefix(ent.name, root) {
+ entsPrinter <- ent
+ rootMet = true
+ } else if rootMet {
+ break
+ }
}
- db.listFiles("", veryFirst)
+ close(entsPrinter)
+ <-printerJob
}