dirsizer.go

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