From 002d4df82f8990e5e7df27c8a35bd7b2e13183f1 Mon Sep 17 00:00:00 2001 From: Felicitas Pojtinger Date: Tue, 30 Nov 2021 23:05:01 +0100 Subject: [PATCH] feat: Add compressionn support to `update` cmd --- cmd/stbak/cmd/archive.go | 60 +++++++++++++++++----------- cmd/stbak/cmd/update.go | 86 ++++++++++++++++++++++++++++++++++------ 2 files changed, 111 insertions(+), 35 deletions(-) diff --git a/cmd/stbak/cmd/archive.go b/cmd/stbak/cmd/archive.go index 9574126..90ee3e4 100644 --- a/cmd/stbak/cmd/archive.go +++ b/cmd/stbak/cmd/archive.go @@ -61,20 +61,7 @@ var archiveCmd = &cobra.Command{ return err } - compressionLevelIsKnown := false - compressionLevel := viper.GetString(compressionLevelFlag) - - for _, candidate := range knownCompressionLevels { - if compressionLevel == candidate { - compressionLevelIsKnown = true - } - } - - if !compressionLevelIsKnown { - return errUnknownCompressionLevel - } - - return nil + return checkCompressionLevel(viper.GetString(compressionLevelFlag)) }, RunE: func(cmd *cobra.Command, args []string) error { if viper.GetBool(verboseFlag) { @@ -234,10 +221,6 @@ func archive( return err } - if err := file.Close(); err != nil { - return err - } - if hdr.PAXRecords == nil { hdr.PAXRecords = map[string]string{} } @@ -302,18 +285,30 @@ func archive( return err } - if err := file.Close(); err != nil { - return err - } - dirty = true return nil }) } +func checkCompressionLevel(compressionLevel string) error { + compressionLevelIsKnown := false + + for _, candidate := range knownCompressionLevels { + if compressionLevel == candidate { + compressionLevelIsKnown = true + } + } + + if !compressionLevelIsKnown { + return errUnknownCompressionLevel + } + + return nil +} + func compress( - file io.Reader, + file io.ReadCloser, tw io.Writer, compressionFormat string, compressionLevel string, @@ -384,6 +379,9 @@ func compress( if err := gz.Close(); err != nil { return err } + if err := file.Close(); err != nil { + return err + } case compressionFormatLZ4Key: l := lz4.Level5 switch compressionLevel { @@ -420,6 +418,9 @@ func compress( if err := lz.Close(); err != nil { return err } + if err := file.Close(); err != nil { + return err + } case compressionFormatZStandardKey: l := zstd.SpeedDefault switch compressionLevel { @@ -459,6 +460,9 @@ func compress( if err := zz.Close(); err != nil { return err } + if err := file.Close(); err != nil { + return err + } case compressionFormatBrotliKey: l := brotli.DefaultCompression switch compressionLevel { @@ -495,6 +499,9 @@ func compress( if err := br.Close(); err != nil { return err } + if err := file.Close(); err != nil { + return err + } case compressionFormatBzip2Key: fallthrough case compressionFormatBzip2ParallelKey: @@ -535,6 +542,9 @@ func compress( if err := bz.Close(); err != nil { return err } + if err := file.Close(); err != nil { + return err + } case compressionFormatNoneKey: if isRegular { if _, err := io.Copy(tw, file); err != nil { @@ -546,6 +556,10 @@ func compress( return err } } + + if err := file.Close(); err != nil { + return err + } default: return errUnsupportedCompressionFormat } diff --git a/cmd/stbak/cmd/update.go b/cmd/stbak/cmd/update.go index 09ea8c5..9cc1695 100644 --- a/cmd/stbak/cmd/update.go +++ b/cmd/stbak/cmd/update.go @@ -3,13 +3,15 @@ package cmd import ( "archive/tar" "context" + "fmt" "io" "io/fs" "os" "path/filepath" + "strconv" "github.com/pojntfx/stfs/pkg/adapters" - "github.com/pojntfx/stfs/pkg/controllers" + "github.com/pojntfx/stfs/pkg/counters" "github.com/pojntfx/stfs/pkg/formatting" "github.com/pojntfx/stfs/pkg/pax" "github.com/pojntfx/stfs/pkg/persisters" @@ -22,6 +24,13 @@ var updateCmd = &cobra.Command{ Use: "update", Aliases: []string{"upd", "u"}, Short: "Update a file or directory's content and metadata on tape or tar file", + PersistentPreRunE: func(cmd *cobra.Command, args []string) error { + if err := viper.BindPFlags(cmd.PersistentFlags()); err != nil { + return err + } + + return checkCompressionLevel(viper.GetString(compressionLevelFlag)) + }, RunE: func(cmd *cobra.Command, args []string) error { if err := viper.BindPFlags(cmd.PersistentFlags()); err != nil { return err @@ -46,6 +55,8 @@ var updateCmd = &cobra.Command{ viper.GetInt(recordSizeFlag), viper.GetString(srcFlag), viper.GetBool(overwriteFlag), + viper.GetString(compressionFlag), + viper.GetString(compressionLevelFlag), ); err != nil { return err } @@ -67,6 +78,8 @@ func update( recordSize int, src string, replacesContent bool, + compressionFormat string, + compressionLevel string, ) error { dirty := false tw, isRegular, cleanup, err := openTapeWriter(tape) @@ -105,6 +118,55 @@ func update( hdr.PAXRecords[pax.STFSRecordVersion] = pax.STFSRecordVersion1 hdr.PAXRecords[pax.STFSRecordAction] = pax.STFSRecordActionUpdate + if info.Mode().IsRegular() && replacesContent { + // Get the compressed size for the header + file, err := os.Open(path) + if err != nil { + return err + } + + fileSizeCounter := counters.CounterWriter{ + Writer: io.Discard, + } + + if err := compress( + file, + &fileSizeCounter, + compressionFormat, + compressionLevel, + isRegular, + recordSize, + ); err != nil { + return err + } + + if hdr.PAXRecords == nil { + hdr.PAXRecords = map[string]string{} + } + hdr.PAXRecords[pax.STFSRecordUncompressedSize] = strconv.Itoa(int(hdr.Size)) + hdr.Size = int64(fileSizeCounter.BytesRead) + + switch compressionFormat { + case compressionFormatGZipKey: + fallthrough + case compressionFormatParallelGZipKey: + hdr.Name += compressionFormatGZipSuffix + case compressionFormatLZ4Key: + hdr.Name += compressionFormatLZ4Suffix + case compressionFormatZStandardKey: + hdr.Name += compressionFormatZStandardSuffix + case compressionFormatBrotliKey: + hdr.Name += compressionFormatBrotliSuffix + case compressionFormatBzip2Key: + fallthrough + case compressionFormatBzip2ParallelKey: + hdr.Name += compressionFormatBzip2Suffix + case compressionFormatNoneKey: + default: + return errUnsupportedCompressionFormat + } + } + if first { if err := formatting.PrintCSV(formatting.TARHeaderCSV); err != nil { return err @@ -128,22 +190,21 @@ func update( return nil } - // TODO: Implement compression + // Compress and write the file file, err := os.Open(path) if err != nil { return err } - defer file.Close() - if isRegular { - if _, err := io.Copy(tw, file); err != nil { - return err - } - } else { - buf := make([]byte, controllers.BlockSize*recordSize) - if _, err := io.CopyBuffer(tw, file, buf); err != nil { - return err - } + if err := compress( + file, + tw, + compressionFormat, + compressionLevel, + isRegular, + recordSize, + ); err != nil { + return err } } else { hdr.Size = 0 // Don't try to seek after the record @@ -167,6 +228,7 @@ func init() { updateCmd.PersistentFlags().IntP(recordSizeFlag, "e", 20, "Amount of 512-bit blocks per record") updateCmd.PersistentFlags().StringP(srcFlag, "s", "", "Path of the file or directory to update") updateCmd.PersistentFlags().BoolP(overwriteFlag, "o", false, "Replace the content on the tape or tar file") + updateCmd.PersistentFlags().StringP(compressionLevelFlag, "l", compressionLevelBalanced, fmt.Sprintf("Compression level to use (default %v, available are %v)", compressionLevelBalanced, knownCompressionLevels)) viper.AutomaticEnv()