github.com/awesome-flow/flow@v0.0.3-0.20190918184116-508d75d68a2c/pkg/corev1alpha1/actor/sink_head_file.go (about)

     1  package actor
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"os"
     7  )
     8  
     9  type FileOpener func(string) (io.WriteCloser, error)
    10  
    11  var DefaultFileOpener = func(path string) (io.WriteCloser, error) {
    12  	switch path {
    13  	case "STDOUT":
    14  		return os.Stdout, nil
    15  	case "STDERR":
    16  		return os.Stderr, nil
    17  	}
    18  	return os.OpenFile(path, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0644)
    19  }
    20  
    21  type SinkHeadFile struct {
    22  	path string
    23  	out  io.WriteCloser
    24  
    25  	Opener FileOpener
    26  }
    27  
    28  var _ (SinkHead) = (*SinkHeadFile)(nil)
    29  
    30  func NewSinkHeadFile(path string) (*SinkHeadFile, error) {
    31  	return &SinkHeadFile{
    32  		path:   path,
    33  		Opener: DefaultFileOpener,
    34  	}, nil
    35  }
    36  
    37  func (h *SinkHeadFile) Connect() error {
    38  	out, err := h.Opener(h.path)
    39  	if err != nil {
    40  		return err
    41  	}
    42  	h.out = out
    43  
    44  	return nil
    45  }
    46  
    47  func (h *SinkHeadFile) Stop() error {
    48  	if h.out != nil {
    49  		if h.out != os.Stdin && h.out != os.Stdout {
    50  			return h.out.Close()
    51  		}
    52  	}
    53  	return nil
    54  }
    55  
    56  func (h *SinkHeadFile) Start() error {
    57  	return nil
    58  }
    59  
    60  func (h *SinkHeadFile) Write(data []byte) (int, error, bool) {
    61  	if h.out == nil {
    62  		return 0, fmt.Errorf("sink head out file is nil"), true
    63  	}
    64  	payload := make([]byte, len(data)+2)
    65  	copy(payload, data)
    66  	copy(payload[len(data):], []byte("\r\n"))
    67  	n, err := h.out.Write(payload)
    68  	if err != nil {
    69  		return 0, err, true
    70  	}
    71  
    72  	return n, nil, false
    73  }