file.go raw

   1  /*
   2   * SPDX-FileCopyrightText: © Hypermode Inc. <hello@hypermode.com>
   3   * SPDX-License-Identifier: Apache-2.0
   4   */
   5  
   6  package z
   7  
   8  import (
   9  	"encoding/binary"
  10  	"errors"
  11  	"fmt"
  12  	"io"
  13  	"os"
  14  	"path/filepath"
  15  )
  16  
  17  // MmapFile represents an mmapd file and includes both the buffer to the data
  18  // and the file descriptor.
  19  type MmapFile struct {
  20  	Data []byte
  21  	Fd   *os.File
  22  }
  23  
  24  var NewFile = errors.New("Create a new file")
  25  
  26  func OpenMmapFileUsing(fd *os.File, sz int, writable bool) (*MmapFile, error) {
  27  	filename := fd.Name()
  28  	fi, err := fd.Stat()
  29  	if err != nil {
  30  		return nil, errors.Join(err, fmt.Errorf("cannot stat file: %s", filename))
  31  	}
  32  
  33  	var rerr error
  34  	fileSize := fi.Size()
  35  	if sz > 0 && fileSize == 0 {
  36  		// If file is empty, truncate it to sz.
  37  		if err := fd.Truncate(int64(sz)); err != nil {
  38  			return nil, errors.Join(err, errors.New("error while truncation"))
  39  		}
  40  		fileSize = int64(sz)
  41  		rerr = NewFile
  42  	}
  43  
  44  	// fmt.Printf("Mmaping file: %s with writable: %v filesize: %d\n", fd.Name(), writable, fileSize)
  45  	buf, err := Mmap(fd, writable, fileSize) // Mmap up to file size.
  46  	if err != nil {
  47  		return nil, errors.Join(err, fmt.Errorf("while mmapping %s with size: %d", fd.Name(), fileSize))
  48  	}
  49  
  50  	if fileSize == 0 {
  51  		dir, _ := filepath.Split(filename)
  52  		if err := SyncDir(dir); err != nil {
  53  			return nil, err
  54  		}
  55  	}
  56  	return &MmapFile{
  57  		Data: buf,
  58  		Fd:   fd,
  59  	}, rerr
  60  }
  61  
  62  // OpenMmapFile opens an existing file or creates a new file. If the file is
  63  // created, it would truncate the file to maxSz. In both cases, it would mmap
  64  // the file to maxSz and returned it. In case the file is created, z.NewFile is
  65  // returned.
  66  func OpenMmapFile(filename string, flag int, maxSz int) (*MmapFile, error) {
  67  	// fmt.Printf("opening file %s with flag: %v\n", filename, flag)
  68  	fd, err := os.OpenFile(filename, flag, 0666)
  69  	if err != nil {
  70  		return nil, errors.Join(err, fmt.Errorf("unable to open: %s", filename))
  71  	}
  72  	writable := true
  73  	if flag == os.O_RDONLY {
  74  		writable = false
  75  	}
  76  	return OpenMmapFileUsing(fd, maxSz, writable)
  77  }
  78  
  79  type mmapReader struct {
  80  	Data   []byte
  81  	offset int
  82  }
  83  
  84  func (mr *mmapReader) Read(buf []byte) (int, error) {
  85  	if mr.offset > len(mr.Data) {
  86  		return 0, io.EOF
  87  	}
  88  	n := copy(buf, mr.Data[mr.offset:])
  89  	mr.offset += n
  90  	if n < len(buf) {
  91  		return n, io.EOF
  92  	}
  93  	return n, nil
  94  }
  95  
  96  func (m *MmapFile) NewReader(offset int) io.Reader {
  97  	return &mmapReader{
  98  		Data:   m.Data,
  99  		offset: offset,
 100  	}
 101  }
 102  
 103  // Bytes returns data starting from offset off of size sz. If there's not enough data, it would
 104  // return nil slice and io.EOF.
 105  func (m *MmapFile) Bytes(off, sz int) ([]byte, error) {
 106  	if len(m.Data[off:]) < sz {
 107  		return nil, io.EOF
 108  	}
 109  	return m.Data[off : off+sz], nil
 110  }
 111  
 112  // Slice returns the slice at the given offset.
 113  func (m *MmapFile) Slice(offset int) []byte {
 114  	sz := binary.BigEndian.Uint32(m.Data[offset:])
 115  	start := offset + 4
 116  	next := start + int(sz)
 117  	if next > len(m.Data) {
 118  		return []byte{}
 119  	}
 120  	res := m.Data[start:next]
 121  	return res
 122  }
 123  
 124  // AllocateSlice allocates a slice of the given size at the given offset.
 125  func (m *MmapFile) AllocateSlice(sz, offset int) ([]byte, int, error) {
 126  	start := offset + 4
 127  
 128  	// If the file is too small, double its size or increase it by 1GB, whichever is smaller.
 129  	if start+sz > len(m.Data) {
 130  		const oneGB = 1 << 30
 131  		growBy := len(m.Data)
 132  		if growBy > oneGB {
 133  			growBy = oneGB
 134  		}
 135  		if growBy < sz+4 {
 136  			growBy = sz + 4
 137  		}
 138  		if err := m.Truncate(int64(len(m.Data) + growBy)); err != nil {
 139  			return nil, 0, err
 140  		}
 141  	}
 142  
 143  	binary.BigEndian.PutUint32(m.Data[offset:], uint32(sz))
 144  	return m.Data[start : start+sz], start + sz, nil
 145  }
 146  
 147  func (m *MmapFile) Sync() error {
 148  	if m == nil {
 149  		return nil
 150  	}
 151  	return Msync(m.Data)
 152  }
 153  
 154  func (m *MmapFile) Delete() error {
 155  	// Badger can set the m.Data directly, without setting any Fd. In that case, this should be a
 156  	// NOOP.
 157  	if m.Fd == nil {
 158  		return nil
 159  	}
 160  
 161  	if err := Munmap(m.Data); err != nil {
 162  		return fmt.Errorf("while munmap file: %s, error: %v\n", m.Fd.Name(), err)
 163  	}
 164  	m.Data = nil
 165  	if err := m.Fd.Truncate(0); err != nil {
 166  		return fmt.Errorf("while truncate file: %s, error: %v\n", m.Fd.Name(), err)
 167  	}
 168  	if err := m.Fd.Close(); err != nil {
 169  		return fmt.Errorf("while close file: %s, error: %v\n", m.Fd.Name(), err)
 170  	}
 171  	return os.Remove(m.Fd.Name())
 172  }
 173  
 174  // Close would close the file. It would also truncate the file if maxSz >= 0.
 175  func (m *MmapFile) Close(maxSz int64) error {
 176  	// Badger can set the m.Data directly, without setting any Fd. In that case, this should be a
 177  	// NOOP.
 178  	if m.Fd == nil {
 179  		return nil
 180  	}
 181  	if err := m.Sync(); err != nil {
 182  		return fmt.Errorf("while sync file: %s, error: %v\n", m.Fd.Name(), err)
 183  	}
 184  	if err := Munmap(m.Data); err != nil {
 185  		return fmt.Errorf("while munmap file: %s, error: %v\n", m.Fd.Name(), err)
 186  	}
 187  	if maxSz >= 0 {
 188  		if err := m.Fd.Truncate(maxSz); err != nil {
 189  			return fmt.Errorf("while truncate file: %s, error: %v\n", m.Fd.Name(), err)
 190  		}
 191  	}
 192  	return m.Fd.Close()
 193  }
 194  
 195  func SyncDir(dir string) error {
 196  	df, err := os.Open(dir)
 197  	if err != nil {
 198  		return errors.Join(err, fmt.Errorf("while opening %s", dir))
 199  	}
 200  	if err := df.Sync(); err != nil {
 201  		return errors.Join(err, fmt.Errorf("while syncing %s", dir))
 202  	}
 203  	if err := df.Close(); err != nil {
 204  		return errors.Join(err, fmt.Errorf("while closing %s", dir))
 205  	}
 206  	return nil
 207  }
 208