dirsizer.go

   1 /*
   2 glocate -- ZFS-diff-friendly locate-like utility
   3 Copyright (C) 2022-2023 Sergey Matveev <stargrave@stargrave.org>
   4
   5 This program is free software: you can redistribute it and/or modify
   6 it under the terms of the GNU General Public License as published by
   7 the Free Software Foundation, version 3 of the License.
   8
   9 This program is distributed in the hope that it will be useful,
  10 but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 GNU General Public License for more details.
  13
  14 You should have received a copy of the GNU General Public License
  15 along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16 */
  17
  18 package main
  19
  20 import (
  21         "bufio"
  22         "encoding/binary"
  23         "io"
  24         "log"
  25         "os"
  26
  27         "github.com/klauspost/compress/zstd"
  28 )
  29
  30 func dirSizer(dirSizes *[]int64, depth int, sinkBack, sinkIn, sinkOut chan Ent) (curSize int64) {
  31         var ent Ent
  32         var opened bool
  33         var dirIdx int
  34         for {
  35                 select {
  36                 case ent = <-sinkBack:
  37                         goto Got
  38                 default:
  39                 }
  40                 ent, opened = <-sinkIn
  41                 if !opened {
  42                         break
  43                 }
  44         Got:
  45                 if len(ent.name) < depth {
  46                         sinkBack <- ent
  47                         return
  48                 }
  49                 sinkOut <- ent
  50                 curSize += ent.size
  51                 if !ent.IsDir() {
  52                         continue
  53                 }
  54                 dirIdx = len(*dirSizes)
  55                 (*dirSizes) = append(*dirSizes, 0)
  56                 dirSize := dirSizer(dirSizes, depth+1, sinkBack, sinkIn, sinkOut)
  57                 (*dirSizes)[dirIdx] = dirSize
  58                 curSize += dirSize
  59         }
  60         return
  61 }
  62
  63 func applyDirSizes(src *os.File, dirSizes []int64) *os.File {
  64         _, err := src.Seek(0, io.SeekStart)
  65         if err != nil {
  66                 log.Fatalln(err)
  67         }
  68         tmp, err := os.CreateTemp(TmpDir, "glocate-idx")
  69         if err != nil {
  70                 log.Fatalln(err)
  71         }
  72
  73         compR, err := zstd.NewReader(src)
  74         if err != nil {
  75                 log.Fatalln(err)
  76         }
  77         br := bufio.NewReaderSize(compR, 1<<17)
  78
  79         compW, err := zstd.NewWriter(tmp,
  80                 zstd.WithEncoderLevel(zstd.SpeedBestCompression))
  81         if err != nil {
  82                 log.Fatalln(err)
  83         }
  84         bw := bufio.NewWriterSize(compW, 1<<17)
  85
  86         num := make([]byte, 8)
  87         var nameLen int
  88         name := make([]byte, 0, 1<<16)
  89         var dirIdx int
  90         for {
  91                 if _, err = io.ReadFull(br, num[:2]); err != nil {
  92                         if err == io.EOF {
  93                                 break
  94                         }
  95                         log.Fatalln(err)
  96                 }
  97                 mustWrite(bw, num[:2])
  98                 nameLen = int(binary.BigEndian.Uint16(num[:2]))
  99                 name = name[:nameLen]
 100                 mustReadFull(br, name)
 101                 mustWrite(bw, name)
 102                 if _, err = io.CopyN(bw, br, 1+8); err != nil {
 103                         log.Fatalln(err)
 104                 }
 105                 if name[len(name)-1] == byte('/') {
 106                         if _, err = br.Discard(8); err != nil {
 107                                 log.Fatalln(err)
 108                         }
 109                         binary.BigEndian.PutUint64(num, uint64(dirSizes[dirIdx]))
 110                         mustWrite(bw, num)
 111                         dirIdx++
 112                 } else {
 113                         if _, err = io.CopyN(bw, br, 8); err != nil {
 114                                 log.Fatalln(err)
 115                         }
 116                 }
 117         }
 118         if err = bw.Flush(); err != nil {
 119                 log.Fatalln(err)
 120         }
 121         if err = compW.Close(); err != nil {
 122                 log.Fatalln(err)
 123         }
 124         compR.Close()
 125         return tmp
 126 }