github.com/fiatjaf/generic-ristretto@v0.0.1/z/file.go (about)

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