github.com/cnotch/ipchub@v1.1.0/av/format/hls/segmentfile.go (about)

     1  // Copyright calabashdad. https://github.com/calabashdad/seal.git
     2  //
     3  // Copyright (c) 2019,CAOHONGJU All rights reserved.
     4  // Use of this source code is governed by a MIT-style
     5  // license that can be found in the LICENSE file.
     6  
     7  package hls
     8  
     9  import (
    10  	"bufio"
    11  	"bytes"
    12  	"fmt"
    13  	"io"
    14  	"os"
    15  	"sync"
    16  
    17  	"github.com/cnotch/ipchub/av/format/mpegts"
    18  )
    19  
    20  type segmentFile interface {
    21  	open(path string) error
    22  	close() error
    23  	writeFrame(frame *mpegts.Frame) error
    24  	get() (io.Reader, int, error)
    25  	delete() error
    26  }
    27  
    28  var segmentPool = sync.Pool{
    29  	New: func() interface{} {
    30  		return bytes.NewBuffer(make([]byte, 0, 512*1024))
    31  	},
    32  }
    33  
    34  type memorySegmentFile struct {
    35  	file *bytes.Buffer
    36  	w    mpegts.FrameWriter
    37  }
    38  
    39  func newMemorySegmentFile() segmentFile {
    40  	return &memorySegmentFile{}
    41  }
    42  
    43  func (mf *memorySegmentFile) open(path string) (err error) {
    44  	mf.file = segmentPool.Get().(*bytes.Buffer)
    45  	mf.file.Reset()
    46  	mf.w, err = mpegts.NewWriter(mf.file)
    47  	return
    48  }
    49  
    50  func (mf *memorySegmentFile) writeFrame(frame *mpegts.Frame) (err error) {
    51  	return mf.w.WriteMpegtsFrame(frame)
    52  }
    53  
    54  func (mf *memorySegmentFile) close() (err error) {
    55  	mf.w = nil
    56  	return
    57  }
    58  
    59  func (mf *memorySegmentFile) get() (io.Reader, int, error) {
    60  	data := mf.file.Bytes()
    61  	return bytes.NewReader(data), len(data), nil
    62  }
    63  
    64  func (mf *memorySegmentFile) delete() error {
    65  	if mf.file != nil {
    66  		segmentPool.Put(mf.file)
    67  		mf.file = nil
    68  	}
    69  	return nil
    70  }
    71  
    72  type persistentSegmentFile struct {
    73  	path string
    74  	file *os.File
    75  	buff *bufio.Writer
    76  	w    mpegts.FrameWriter
    77  }
    78  
    79  func newPersistentSegmentFile() segmentFile {
    80  	return &persistentSegmentFile{}
    81  }
    82  
    83  func (pf *persistentSegmentFile) open(path string) (err error) {
    84  	pf.path = path
    85  	pf.file, err = os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, os.ModePerm)
    86  	if err != nil {
    87  		return fmt.Errorf("open file error, file=%s", path)
    88  	}
    89  
    90  	pf.buff = bufio.NewWriterSize(pf.file, 64*1024)
    91  	pf.w, err = mpegts.NewWriter(pf.buff)
    92  	return
    93  }
    94  
    95  func (pf *persistentSegmentFile) writeFrame(frame *mpegts.Frame) (err error) {
    96  	return pf.w.WriteMpegtsFrame(frame)
    97  }
    98  
    99  func (pf *persistentSegmentFile) close() (err error) {
   100  	if nil == pf.file {
   101  		return
   102  	}
   103  
   104  	pf.buff.Flush()
   105  	pf.file.Close()
   106  
   107  	// after close, rest the file write to nil
   108  	pf.file = nil
   109  	pf.buff = nil
   110  	pf.w = nil
   111  	return nil
   112  }
   113  
   114  func (pf *persistentSegmentFile) get() (reader io.Reader, size int, err error) {
   115  	var finfo os.FileInfo
   116  	finfo, err = os.Stat(pf.path)
   117  	if err != nil {
   118  		return
   119  	}
   120  	var f *os.File
   121  	if f, err = os.Open(pf.path); err != nil {
   122  		return nil, 0, err
   123  	}
   124  
   125  	return f, int(finfo.Size()), nil
   126  }
   127  
   128  func (pf *persistentSegmentFile) delete() error {
   129  	pf.close()
   130  	return os.Remove(pf.path)
   131  }