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