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

     1  package actor
     2  
     3  import (
     4  	"bytes"
     5  	"compress/flate"
     6  	"compress/gzip"
     7  	"compress/lzw"
     8  	"compress/zlib"
     9  	"fmt"
    10  	"sync"
    11  
    12  	"github.com/DataDog/zstd"
    13  	core "github.com/awesome-flow/flow/pkg/corev1alpha1"
    14  	"github.com/golang/snappy"
    15  )
    16  
    17  type CoderFunc func([]byte, int) ([]byte, error)
    18  
    19  var DefaultCoders = map[string]CoderFunc{
    20  	"gzip": func(payload []byte, level int) ([]byte, error) {
    21  		var b bytes.Buffer
    22  		w, err := gzip.NewWriterLevel(&b, level)
    23  		if err != nil {
    24  			return nil, err
    25  		}
    26  		if _, err := w.Write(payload); err != nil {
    27  			return nil, err
    28  		}
    29  		w.Close()
    30  		return b.Bytes(), nil
    31  	},
    32  	"flate": func(payload []byte, level int) ([]byte, error) {
    33  		var b bytes.Buffer
    34  		w, err := flate.NewWriter(&b, level)
    35  		if err != nil {
    36  			return nil, err
    37  		}
    38  		if _, err := w.Write(payload); err != nil {
    39  			return nil, err
    40  		}
    41  		w.Close()
    42  		return b.Bytes(), nil
    43  	},
    44  	"lzw": func(payload []byte, _ int) ([]byte, error) {
    45  		var b bytes.Buffer
    46  		// The final digit is the literal coder width. Varies from 2 to
    47  		// 8 bits. We are using 8 by default here.
    48  		// See https://golang.org/src/compress/lzw/writer.go#L241
    49  		// for more details.
    50  		w := lzw.NewWriter(&b, lzw.MSB, 8)
    51  		if _, err := w.Write(payload); err != nil {
    52  			return nil, err
    53  		}
    54  		w.Close()
    55  		return b.Bytes(), nil
    56  	},
    57  	"zlib": func(payload []byte, level int) ([]byte, error) {
    58  		var b bytes.Buffer
    59  		w, err := zlib.NewWriterLevel(&b, level)
    60  		if err != nil {
    61  			return nil, err
    62  		}
    63  		if _, err := w.Write(payload); err != nil {
    64  			return nil, err
    65  		}
    66  		w.Close()
    67  		return b.Bytes(), nil
    68  	},
    69  	"zstd": func(payload []byte, level int) ([]byte, error) {
    70  		var b bytes.Buffer
    71  		w := zstd.NewWriterLevel(&b, level)
    72  		if _, err := w.Write(payload); err != nil {
    73  			return nil, err
    74  		}
    75  		w.Close()
    76  		return b.Bytes(), nil
    77  	},
    78  	"snappy": func(payload []byte, _ int) ([]byte, error) {
    79  		var b bytes.Buffer
    80  		w := snappy.NewBufferedWriter(&b)
    81  		if _, err := w.Write(payload); err != nil {
    82  			return nil, err
    83  		}
    84  		w.Close()
    85  		return b.Bytes(), nil
    86  	},
    87  }
    88  
    89  type Compressor struct {
    90  	name  string
    91  	ctx   *core.Context
    92  	coder CoderFunc
    93  	level int
    94  	queue chan *core.Message
    95  	wg    sync.WaitGroup
    96  }
    97  
    98  var _ core.Actor = (*Compressor)(nil)
    99  
   100  func NewCompressor(name string, ctx *core.Context, params core.Params) (core.Actor, error) {
   101  	return NewCompressorWithCoders(name, ctx, params, DefaultCoders)
   102  }
   103  
   104  func NewCompressorWithCoders(name string, ctx *core.Context, params core.Params, coders map[string]CoderFunc) (core.Actor, error) {
   105  	alg, ok := params["compress"]
   106  	if !ok {
   107  		return nil, fmt.Errorf("compressor %q is missing `compress` config", name)
   108  	}
   109  	coder, ok := coders[alg.(string)]
   110  	if !ok {
   111  		return nil, fmt.Errorf("compressor %q: unknown compression algorithm %q", name, alg)
   112  	}
   113  	level := -1
   114  	if l, ok := params["level"]; ok {
   115  		if _, ok := l.(int); !ok {
   116  			return nil, fmt.Errorf("compressor %q: malformed compression level provided: got: %+v, want: an integer", name, l)
   117  		}
   118  		level = l.(int)
   119  	}
   120  
   121  	return &Compressor{
   122  		name:  name,
   123  		ctx:   ctx,
   124  		coder: coder,
   125  		level: level,
   126  		queue: make(chan *core.Message),
   127  	}, nil
   128  }
   129  
   130  func (c *Compressor) Name() string {
   131  	return c.name
   132  }
   133  
   134  func (c *Compressor) Start() error {
   135  	return nil
   136  }
   137  
   138  func (c *Compressor) Stop() error {
   139  	close(c.queue)
   140  	c.wg.Wait()
   141  
   142  	return nil
   143  }
   144  
   145  func (c *Compressor) Connect(nthreads int, peer core.Receiver) error {
   146  	for i := 0; i < nthreads; i++ {
   147  		c.wg.Add(1)
   148  		go func() {
   149  			for msg := range c.queue {
   150  				if err := peer.Receive(msg); err != nil {
   151  					msg.Complete(core.MsgStatusFailed)
   152  					c.ctx.Logger().Error(err.Error())
   153  				}
   154  			}
   155  			c.wg.Done()
   156  		}()
   157  	}
   158  
   159  	return nil
   160  }
   161  
   162  func (c *Compressor) Receive(msg *core.Message) error {
   163  	data, err := c.coder(msg.Body(), c.level)
   164  	if err != nil {
   165  		return err
   166  	}
   167  	msg.SetBody(data)
   168  	c.queue <- msg
   169  
   170  	return nil
   171  }