github.com/whiteboxio/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 }