github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/ext/dsort/shard/write_compressed.go (about)

     1  // Package shard provides Extract(shard), Create(shard), and associated methods
     2  // across all suppported archival formats (see cmn/archive/mime.go)
     3  /*
     4   * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
     5   */
     6  package shard
     7  
     8  import (
     9  	"archive/tar"
    10  	"io"
    11  
    12  	"github.com/NVIDIA/aistore/cmn/archive"
    13  	"github.com/NVIDIA/aistore/cmn/cos"
    14  	"github.com/NVIDIA/aistore/cmn/debug"
    15  )
    16  
    17  // common method to compress .tar via `aw` (archive writer)
    18  func writeCompressedTar(s *Shard, tw *tar.Writer, aw io.Writer, loader ContentLoader, rdReader *tarRecordW) (written int64, _ error) {
    19  	var needFlush bool
    20  	for _, rec := range s.Records.All() {
    21  		for _, obj := range rec.Objects {
    22  			switch obj.StoreType {
    23  			case OffsetStoreType:
    24  				if needFlush {
    25  					// write directly to the underlying tar file, so flush everything written so far
    26  					if err := tw.Flush(); err != nil {
    27  						return written, err
    28  					}
    29  					needFlush = false
    30  				}
    31  				n, err := loader.Load(aw, rec, obj)
    32  				written += n
    33  				if err != nil {
    34  					return written, err
    35  				}
    36  				// pad to 512
    37  				diff := cos.CeilAlignInt64(n, archive.TarBlockSize) - n
    38  				debug.Assert(diff >= 0)
    39  				if diff > 0 {
    40  					npad, errP := aw.Write(padBuf[:diff])
    41  					written += int64(npad)
    42  					if errP != nil {
    43  						return written, errP
    44  					}
    45  					debug.Assert(diff == int64(npad) && diff < archive.TarBlockSize, diff, npad)
    46  				}
    47  			case SGLStoreType, DiskStoreType:
    48  				rdReader.reinit(tw, obj.Size, obj.MetadataSize)
    49  				n, err := loader.Load(rdReader, rec, obj)
    50  				written += n
    51  				if err != nil {
    52  					return written, err
    53  				}
    54  				needFlush = true
    55  			default:
    56  				debug.Assert(false, obj.StoreType)
    57  			}
    58  		}
    59  	}
    60  	return written, nil
    61  }