github.com/outcaste-io/ristretto@v0.2.3/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 }