github.com/sohaha/zlsgo@v1.7.13-0.20240501141223-10dd1a906f76/zfile/memory.go (about)

     1  package zfile
     2  
     3  import (
     4  	"io"
     5  	"os"
     6  	"sync"
     7  	"time"
     8  )
     9  
    10  type MemoryFile struct {
    11  	stopTiming chan struct{}
    12  	fbefore    memoryFileFlushBefore
    13  	name       string
    14  	buffer     byteBuffer
    15  	timing     int64
    16  	lock       sync.RWMutex
    17  	stop       bool
    18  }
    19  
    20  type MemoryFileOption func(*MemoryFile)
    21  type memoryFileFlushBefore func(f *MemoryFile) error
    22  
    23  func MemoryFileAutoFlush(second int64) func(*MemoryFile) {
    24  	return func(f *MemoryFile) {
    25  		f.timing = second
    26  	}
    27  }
    28  
    29  // MemoryFileFlushBefore is a function that will be called before flush to disk,take care to avoid writing to prevent deadlocks
    30  func MemoryFileFlushBefore(fn memoryFileFlushBefore) func(*MemoryFile) {
    31  	return func(f *MemoryFile) {
    32  		f.fbefore = fn
    33  	}
    34  }
    35  
    36  func NewMemoryFile(name string, opt ...MemoryFileOption) *MemoryFile {
    37  	f := &MemoryFile{name: name, buffer: makeByteBuffer([]byte{})}
    38  	for _, o := range opt {
    39  		o(f)
    40  	}
    41  	f.stopTiming = make(chan struct{})
    42  	if f.timing > 0 {
    43  		go f.flushLoop()
    44  	}
    45  	return f
    46  }
    47  
    48  func (f *MemoryFile) flushLoop() {
    49  	ticker := time.NewTicker(time.Second * time.Duration(f.timing))
    50  	for {
    51  		select {
    52  		case <-f.stopTiming:
    53  			return
    54  		case <-ticker.C:
    55  			_ = f.Sync()
    56  		}
    57  	}
    58  }
    59  
    60  func (f *MemoryFile) SetName(name string) {
    61  	f.name = name
    62  }
    63  
    64  func (f *MemoryFile) Bytes() []byte {
    65  	f.lock.RLock()
    66  	b := f.buffer.buffer
    67  	f.lock.RUnlock()
    68  	return b
    69  }
    70  
    71  func (f *MemoryFile) Stat() (os.FileInfo, error) {
    72  	return f, nil
    73  }
    74  
    75  func (f *MemoryFile) Read(buffer []byte) (int, error) {
    76  	f.lock.RLock()
    77  	n, err := f.buffer.Read(buffer)
    78  	f.lock.RUnlock()
    79  	return n, err
    80  }
    81  
    82  func (f *MemoryFile) Close() error {
    83  	f.lock.Lock()
    84  	if f.stop {
    85  		return nil
    86  	}
    87  	f.stop = true
    88  	f.lock.Unlock()
    89  	f.stopTiming <- struct{}{}
    90  	return f.Sync()
    91  }
    92  
    93  func (f *MemoryFile) Sync() error {
    94  	if f.Size() == 0 {
    95  		return nil
    96  	}
    97  	if f.fbefore != nil {
    98  		err := f.fbefore(f)
    99  		if err != nil {
   100  			return err
   101  		}
   102  	}
   103  	f.lock.Lock()
   104  	b := f.buffer.buffer
   105  	f.buffer.Reset()
   106  	f.lock.Unlock()
   107  	return WriteFile(f.name, b, true)
   108  }
   109  
   110  func (f *MemoryFile) Write(buffer []byte) (int, error) {
   111  	f.lock.Lock()
   112  	n, err := f.buffer.Write(buffer)
   113  	f.lock.Unlock()
   114  	return n, err
   115  }
   116  
   117  func (f *MemoryFile) Seek(offset int64, whence int) (int64, error) {
   118  	f.lock.RLock()
   119  	n, err := f.buffer.Seek(offset, whence)
   120  	f.lock.RUnlock()
   121  	return n, err
   122  }
   123  
   124  func (f *MemoryFile) Name() string {
   125  	return f.name
   126  }
   127  
   128  func (f *MemoryFile) Size() int64 {
   129  	f.lock.RLock()
   130  	l := int64(f.buffer.Len())
   131  	f.lock.RUnlock()
   132  	return l
   133  }
   134  
   135  func (f *MemoryFile) Mode() os.FileMode {
   136  	return 0666
   137  }
   138  
   139  func (f *MemoryFile) ModTime() time.Time {
   140  	return time.Time{}
   141  }
   142  
   143  func (f *MemoryFile) IsDir() bool {
   144  	return false
   145  }
   146  
   147  func (f *MemoryFile) Sys() interface{} {
   148  	return nil
   149  }
   150  
   151  type byteBuffer struct {
   152  	buffer []byte
   153  	index  int
   154  }
   155  
   156  func makeByteBuffer(buffer []byte) byteBuffer {
   157  	return byteBuffer{
   158  		buffer: buffer,
   159  		index:  0,
   160  	}
   161  }
   162  
   163  func (bb *byteBuffer) Reset() {
   164  	bb.buffer = bb.buffer[:0]
   165  	bb.index = 0
   166  }
   167  
   168  func (bb *byteBuffer) Len() int {
   169  	return len(bb.buffer)
   170  }
   171  
   172  func (bb *byteBuffer) Position() int {
   173  	return bb.index
   174  }
   175  
   176  func (bb *byteBuffer) Bytes() []byte {
   177  	return bb.buffer
   178  }
   179  
   180  func (bb *byteBuffer) Read(buffer []byte) (int, error) {
   181  	if len(buffer) == 0 {
   182  		return 0, nil
   183  	}
   184  
   185  	if bb.index >= bb.Len() {
   186  		return 0, io.EOF
   187  	}
   188  
   189  	last := copy(buffer, bb.buffer[bb.index:])
   190  	bb.index += last
   191  	return last, nil
   192  }
   193  
   194  func (bb *byteBuffer) Write(buffer []byte) (int, error) {
   195  	bb.buffer = append(bb.buffer[:bb.index], buffer...)
   196  	bb.index += len(buffer)
   197  	return len(buffer), nil
   198  }
   199  
   200  func (bb *byteBuffer) Seek(offset int64, whence int) (int64, error) {
   201  	switch whence {
   202  	default:
   203  	case io.SeekStart:
   204  		bb.index = int(offset)
   205  	case io.SeekCurrent:
   206  		bb.index += int(offset)
   207  	case io.SeekEnd:
   208  		bb.index = bb.Len() - 1 - int(offset)
   209  	}
   210  	return int64(bb.index), nil
   211  }