├── .travis.yml ├── LICENSE ├── README.md ├── compressing_blob.go ├── compressing_blob_test.go ├── entrypoints ├── build_fluentd_forwarder │ └── main.go ├── dummy_td_server │ └── main.go └── fluentd_forwarder │ ├── main.go │ └── signal.go ├── errors.go ├── file_journal.go ├── file_journal_test.go ├── forwarder.go ├── input.go ├── output.go ├── output_td.go ├── output_td_test.go ├── path_builder.go ├── path_builder_test.go ├── utils.go └── worker_set.go /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | 3 | go: 4 | - 1.4 5 | 6 | before_install: 7 | - go get github.com/ugorji/go/codec 8 | - go get github.com/op/go-logging 9 | - go get github.com/jehiah/go-strftime 10 | - go get github.com/moriyoshi/go-ioextras 11 | - go get gopkg.in/gcfg.v1 12 | - go get github.com/treasure-data/td-client-go 13 | 14 | script: 15 | - cd entrypoints/fluentd_forwarder && go build 16 | - cd - && go test 17 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright [yyyy] [name of copyright owner] 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. 203 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | fluentd-forwarder 2 | ================= 3 | 4 | A lightweight Fluentd forwarder written in Go. 5 | 6 | Requirements 7 | ------------ 8 | 9 | - Go v1.4.1 or above 10 | - Set the $GOPATH environment variable to get `fluentd_forwarder` 11 | under `$GOPATH/bin` directory. 12 | 13 | Build Instructions 14 | ------------------ 15 | 16 | To install the required dependencies and build `fluentd_forwarder` do: 17 | 18 | ``` 19 | $ go get github.com/fluent/fluentd-forwarder/entrypoints/build_fluentd_forwarder 20 | $ bin/build_fluentd_forwarder fluentd_forwarder 21 | ``` 22 | 23 | Running `fluentd_forwarder` 24 | --------------------------- 25 | 26 | ``` 27 | $ $GOPATH/bin/fluentd_forwarder 28 | ``` 29 | 30 | Without arguments, it simply listens on 127.0.0.1:24224 and forwards the events to 127.0.0.1:24225. 31 | 32 | It gracefully stops in response to SIGINT. 33 | 34 | If you want to specify where to forward the events, try the following: 35 | 36 | ``` 37 | $ $GOPATH/bin/fluentd_forwarder -to fluent://some-remote-node.local:24224 38 | ``` 39 | 40 | Command-line Options 41 | -------------------- 42 | 43 | * -retry-interval 44 | 45 | Retry interval in which connection is tried against the remote agent. 46 | 47 | ``` 48 | -retry-interval 5s 49 | ``` 50 | 51 | * -conn-timeout 52 | 53 | Connection timeout after which the connection has failed. 54 | 55 | ``` 56 | -conn-timeout 10s 57 | ``` 58 | 59 | * -write-timeout 60 | 61 | Write timeout on wire. 62 | 63 | ``` 64 | -write-timeout 30s 65 | ``` 66 | 67 | * -flush-interval 68 | 69 | Flush interval in which the events are forwareded to the remote agent . 70 | 71 | ``` 72 | -flush-interval 5s 73 | ``` 74 | 75 | * -listen-on 76 | 77 | Interface address and port on which the forwarder listens. 78 | 79 | ``` 80 | -listen-on 127.0.0.1:24224 81 | ``` 82 | 83 | * -to 84 | 85 | Host and port to which the events are forwarded. 86 | 87 | ``` 88 | -to remote-host.local:24225 89 | -to fluent://remote-host.local:24225 90 | -to td+https://urlencoded-api-key@/*/* 91 | -to td+https://urlencoded-api-key@/database/* 92 | -to td+https://urlencoded-api-key@/database/table 93 | -to td+https://urlencoded-api-key@endpoint/*/* 94 | ``` 95 | 96 | * -ca-certs 97 | 98 | SSL CA certficates to be verified against when the secure connection is used. Must be in PEM format. You can use the [one bundled with td-client-ruby](https://raw.githubusercontent.com/treasure-data/td-client-ruby/master/data/ca-bundle.crt). 99 | 100 | ``` 101 | -ca-certs ca-bundle.crt 102 | ``` 103 | 104 | 105 | * -buffer-path 106 | 107 | Directory / path on which buffer files are created. * may be used within the path to indicate the prefix or suffix like var/pre*suf 108 | 109 | ``` 110 | -buffer-path /var/lib/fluent-forwarder 111 | -buffer-path /var/lib/fluent-forwarder/prefix*suffix 112 | ``` 113 | 114 | * -buffer-chunk-limit 115 | 116 | Maximum size of a buffer chunk 117 | 118 | ``` 119 | -buffer-chunk-limit 16777216 120 | ``` 121 | 122 | * -parallelism 123 | 124 | Number of simultaneous connections used to submit events. It takes effect only when the target is td+http(s). 125 | 126 | ``` 127 | -parallelism 1 128 | ``` 129 | 130 | * -log-level 131 | 132 | Logging level. Any one of the following values; CRITICAL, ERROR, WARNING, NOTICE, INFO and DEBUG. 133 | 134 | ``` 135 | -log-level DEBUG 136 | ``` 137 | 138 | * -log-file 139 | 140 | Species the path to the log file. By default logging is performed to the standard error. It may contain strftime(3)-like format specifications like %Y in any positions. If the parent directory doesn't exist at the time the logging is performed, all the leading directories are created automatically so you can specify the path like `/var/log/fluentd_forwarder/%Y-%m-%d/fluentd_forwarder.log` 141 | 142 | ``` 143 | -log-file /var/log/fluentd_forwarder.log 144 | ``` 145 | 146 | * -config 147 | 148 | Specifies the path to the configuration file. The syntax is detailed below. 149 | 150 | ``` 151 | -config /etc/fluentd-forwarder/fluentd-forwarder.cfg 152 | ``` 153 | 154 | * -metadata 155 | 156 | Specifies the additional data to insert `metadata` record. The syntax is detailed below. 157 | ``` 158 | -metadata "custom metadata" 159 | ``` 160 | 161 | Configuration File 162 | ------------------ 163 | 164 | The syntax of the configuration file is so-called INI format with the name of the primary section being `fluentd-forwarder`. Each setting is named exactly the same as the command-line counterpart, except for `-config`. (It is not possible to refer to another configuation file from a configuration file) 165 | 166 | ``` 167 | [fluentd-forwarder] 168 | to = fluent://remote.local:24224 169 | buffer-chunk-limit = 16777216 170 | flush-interval = 10s 171 | retry-interval = 1s 172 | ``` 173 | 174 | Dependencies 175 | ------------ 176 | 177 | fluentd_forwarder depends on the following external libraries: 178 | 179 | * github.com/ugorji/go/codec 180 | * github.com/op/go-logging 181 | * github.com/jehiah/go-strftime 182 | * github.com/moriyoshi/go-ioextras 183 | * gopkg.in/gcfg.v1 184 | * github.com/treasure-data/td-client-go 185 | 186 | License 187 | ------- 188 | 189 | The source code and its object form ("Work"), unless otherwise specified, are licensed under the Apache Software License, Version 2.0. You may not use the Work except in compliance with the License. You may obtain a copy of the License at 190 | 191 | http://www.apache.org/licenses/LICENSE-2.0 192 | 193 | Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. 194 | 195 | A portion of the code originally written by Moriyoshi Koizumi and later modified by Treasure Data, Inc. continues to be published and distributed under the same terms and conditions as the MIT license, with its authorship being attributed to the both parties. It is specified at the top of the applicable source files. 196 | -------------------------------------------------------------------------------- /compressing_blob.go: -------------------------------------------------------------------------------- 1 | // 2 | // Fluentd Forwarder 3 | // 4 | // Copyright (C) 2014 Treasure Data, Inc. 5 | // 6 | // Licensed under the Apache License, Version 2.0 (the "License"); 7 | // you may not use this file except in compliance with the License. 8 | // You may obtain a copy of the License at 9 | // 10 | // http://www.apache.org/licenses/LICENSE-2.0 11 | // 12 | // Unless required by applicable law or agreed to in writing, software 13 | // distributed under the License is distributed on an "AS IS" BASIS, 14 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | // See the License for the specific language governing permissions and 16 | // limitations under the License. 17 | // 18 | 19 | package fluentd_forwarder 20 | 21 | import ( 22 | "bufio" 23 | "compress/gzip" 24 | "crypto/md5" 25 | "errors" 26 | ioextras "github.com/moriyoshi/go-ioextras" 27 | td_client "github.com/treasure-data/td-client-go" 28 | "hash" 29 | "io" 30 | "os" 31 | ) 32 | 33 | type CompressingBlob struct { 34 | inner td_client.Blob 35 | level int 36 | bufferSize int 37 | reader *CompressingBlobReader 38 | tempFactory ioextras.RandomAccessStoreFactory 39 | md5sum []byte 40 | size int64 41 | } 42 | 43 | type CompressingBlobReader struct { 44 | buf []byte // ring buffer 45 | o int 46 | src io.ReadCloser 47 | dst *ioextras.StoreReadWriter 48 | s ioextras.SizedRandomAccessStore 49 | w *ioextras.StoreReadWriter 50 | bw *bufio.Writer 51 | cw *gzip.Writer 52 | h hash.Hash 53 | eof bool 54 | md5SumAvailable bool 55 | closeNotify func(*CompressingBlobReader) 56 | } 57 | 58 | func (reader *CompressingBlobReader) drainAll() error { 59 | rn := 0 60 | err := (error)(nil) 61 | for reader.cw != nil && err == nil { 62 | o := reader.o 63 | if !reader.eof { 64 | rn, err = reader.src.Read(reader.buf[reader.o:cap(reader.buf)]) 65 | o += rn 66 | if err == io.EOF { 67 | reader.eof = true 68 | } 69 | } else { 70 | if o == 0 { 71 | break 72 | } 73 | } 74 | if o > 0 { 75 | wn, werr := reader.cw.Write(reader.buf[0:o]) 76 | if werr != nil { 77 | err = werr 78 | } 79 | copy(reader.buf[0:], reader.buf[wn:]) 80 | reader.o = o - wn 81 | } 82 | if err != nil { 83 | reader.cw.Close() 84 | reader.cw = nil 85 | werr := reader.bw.Flush() 86 | if werr != nil { 87 | return werr 88 | } 89 | reader.bw = nil 90 | } 91 | } 92 | if err == io.EOF { 93 | err = nil 94 | } 95 | return err 96 | } 97 | 98 | func (reader *CompressingBlobReader) Read(p []byte) (int, error) { 99 | n := int(0) 100 | err := (error)(nil) 101 | rn := int(0) 102 | rpos, err := reader.dst.Seek(0, os.SEEK_CUR) 103 | if err != nil { 104 | return 0, err // should never happen 105 | } 106 | wpos, err := reader.w.Seek(0, os.SEEK_CUR) 107 | if err != nil { 108 | return 0, err // should never happen 109 | } 110 | for { 111 | n = int(wpos - rpos) // XXX: may underflow, but it should be ok 112 | if err != nil || len(p) <= n || reader.cw == nil { 113 | break 114 | } 115 | o := reader.o 116 | if !reader.eof { 117 | rn, err = reader.src.Read(reader.buf[reader.o:cap(reader.buf)]) 118 | o += rn 119 | if err == io.EOF { 120 | reader.eof = true 121 | } 122 | } 123 | if o > 0 { 124 | wn, werr := reader.cw.Write(reader.buf[0:o]) 125 | copy(reader.buf[0:], reader.buf[wn:o]) 126 | reader.o = o - wn 127 | if werr != nil { 128 | return 0, werr 129 | } 130 | } else { 131 | if reader.eof && reader.cw != nil { 132 | reader.cw.Close() 133 | reader.cw = nil 134 | werr := reader.bw.Flush() 135 | if werr != nil { 136 | return 0, werr 137 | } 138 | reader.bw = nil 139 | } 140 | } 141 | var werr error 142 | wpos, werr = reader.w.Seek(0, os.SEEK_CUR) 143 | if werr != nil { 144 | return 0, werr // should never happen 145 | } 146 | } 147 | if len(p) > 0 && n == 0 && reader.eof { 148 | reader.md5SumAvailable = true 149 | return 0, io.EOF 150 | } 151 | if n > len(p) { 152 | n = len(p) 153 | } 154 | n, err = reader.dst.Read(p[0:n]) 155 | reader.h.Write(p[0:n]) 156 | if err == io.EOF { 157 | if !reader.eof { 158 | panic("something went wrong!") 159 | } 160 | reader.md5SumAvailable = true 161 | } 162 | return n, err 163 | } 164 | 165 | func (reader *CompressingBlobReader) size() (int64, error) { 166 | if reader.s == nil { 167 | return -1, errors.New("already closed") 168 | } 169 | if reader.cw != nil { 170 | err := reader.drainAll() 171 | if err != nil { 172 | return -1, err 173 | } 174 | } 175 | size, err := reader.s.Size() 176 | if err != nil { 177 | return -1, err // should never happen 178 | } 179 | return size, nil 180 | } 181 | 182 | func (reader *CompressingBlobReader) Close() error { 183 | bwerr := (error)(nil) 184 | errs := make([]error, 0, 4) 185 | err := reader.ensureMD5SumAvailble() 186 | if err != nil { 187 | return err 188 | } 189 | if reader.cw != nil { 190 | if err != nil { 191 | return err 192 | } 193 | err = reader.cw.Close() 194 | if err == nil { 195 | reader.cw = nil 196 | } else { 197 | errs = append(errs, err) 198 | } 199 | } 200 | if reader.bw != nil { 201 | bwerr = reader.bw.Flush() 202 | if bwerr == nil { 203 | reader.bw = nil 204 | } else { 205 | errs = append(errs, bwerr) 206 | } 207 | } 208 | if reader.src != nil { 209 | err := reader.src.Close() 210 | if err == nil { 211 | reader.src = nil 212 | } else { 213 | errs = append(errs, err) 214 | } 215 | } 216 | if bwerr == nil { 217 | if reader.s != nil { 218 | err := reader.s.Close() 219 | if err == nil { 220 | reader.s = nil 221 | } else { 222 | errs = append(errs, err) 223 | } 224 | } 225 | } 226 | if len(errs) > 0 { 227 | return Errors(errs) 228 | } else { 229 | reader.closeNotify(reader) 230 | return nil 231 | } 232 | } 233 | 234 | func (reader *CompressingBlobReader) ensureMD5SumAvailble() error { 235 | if reader.md5SumAvailable { 236 | return nil 237 | } 238 | if reader.s == nil { 239 | return errors.New("already closed") 240 | } 241 | err := reader.drainAll() 242 | if err != nil { 243 | return err 244 | } 245 | r := *reader.dst 246 | _, err = io.Copy(reader.h, &r) 247 | if err != nil { 248 | return err 249 | } 250 | reader.md5SumAvailable = true 251 | return nil 252 | } 253 | 254 | func (reader *CompressingBlobReader) md5sum() ([]byte, error) { 255 | err := reader.ensureMD5SumAvailble() 256 | if err != nil { 257 | return nil, err 258 | } 259 | retval := make([]byte, 0, reader.h.Size()) 260 | return reader.h.Sum(retval), nil 261 | } 262 | 263 | func (blob *CompressingBlob) newReader() (*CompressingBlobReader, error) { 264 | err := (error)(nil) 265 | src := (io.ReadCloser)(nil) 266 | s := (ioextras.SizedRandomAccessStore)(nil) 267 | defer func() { 268 | if err != nil { 269 | if src != nil { 270 | src.Close() 271 | } 272 | if s != nil { 273 | s.Close() 274 | } 275 | } 276 | }() 277 | src, err = blob.inner.Reader() 278 | if err != nil { 279 | return nil, err 280 | } 281 | s_, err := blob.tempFactory.RandomAccessStore() 282 | if err != nil { 283 | return nil, err 284 | } 285 | s = s_.(ioextras.SizedRandomAccessStore) 286 | w := &ioextras.StoreReadWriter{s, 0, -1} 287 | dst := &ioextras.StoreReadWriter{s, 0, -1} 288 | // assuming average compression ratio to be 1/3 289 | writeBufferSize := maxInt(4096, blob.bufferSize/3) 290 | bw := bufio.NewWriterSize(w, writeBufferSize) 291 | cw, err := gzip.NewWriterLevel(bw, blob.level) 292 | if err != nil { 293 | return nil, err 294 | } 295 | return &CompressingBlobReader{ 296 | buf: make([]byte, blob.bufferSize), 297 | o: 0, 298 | src: src, 299 | dst: dst, 300 | s: s, 301 | w: w, 302 | bw: bw, 303 | cw: cw, 304 | eof: false, 305 | h: md5.New(), 306 | md5SumAvailable: false, 307 | closeNotify: func(reader *CompressingBlobReader) { 308 | md5sum, err := reader.md5sum() 309 | if err == nil { 310 | blob.md5sum = md5sum 311 | } 312 | size, err := reader.size() 313 | if err == nil { 314 | blob.size = size 315 | } 316 | blob.reader = nil 317 | }, 318 | }, nil 319 | } 320 | 321 | func (blob *CompressingBlob) ensureReaderAvailable() error { 322 | if blob.reader != nil { 323 | return nil 324 | } 325 | reader, err := blob.newReader() 326 | if err != nil { 327 | return err 328 | } 329 | blob.reader = reader 330 | blob.md5sum = nil 331 | blob.size = -1 332 | return nil 333 | } 334 | 335 | func (blob *CompressingBlob) Reader() (io.ReadCloser, error) { 336 | err := blob.ensureReaderAvailable() 337 | if err != nil { 338 | return nil, err 339 | } 340 | return blob.reader, nil 341 | } 342 | 343 | func (blob *CompressingBlob) Size() (int64, error) { 344 | if blob.size < 0 { 345 | err := blob.ensureReaderAvailable() 346 | if err != nil { 347 | return -1, err 348 | } 349 | size, err := blob.reader.size() 350 | if err != nil { 351 | return -1, err 352 | } 353 | blob.size = size 354 | } 355 | return blob.size, nil 356 | } 357 | 358 | func (blob *CompressingBlob) MD5Sum() ([]byte, error) { 359 | if blob.md5sum == nil { 360 | err := blob.ensureReaderAvailable() 361 | if err != nil { 362 | return nil, err 363 | } 364 | md5sum, err := blob.reader.md5sum() 365 | if err != nil { 366 | return nil, err 367 | } 368 | blob.md5sum = md5sum 369 | } 370 | return blob.md5sum, nil 371 | } 372 | 373 | func (blob *CompressingBlob) Dispose() error { 374 | if blob.reader != nil { 375 | return blob.reader.Close() 376 | } 377 | return nil 378 | } 379 | 380 | func NewCompressingBlob(blob td_client.Blob, bufferSize int, level int, tempFactory ioextras.RandomAccessStoreFactory) *CompressingBlob { 381 | return &CompressingBlob{ 382 | inner: blob, 383 | level: level, 384 | bufferSize: bufferSize, 385 | reader: nil, 386 | tempFactory: tempFactory, 387 | md5sum: nil, 388 | size: -1, 389 | } 390 | } 391 | -------------------------------------------------------------------------------- /compressing_blob_test.go: -------------------------------------------------------------------------------- 1 | // 2 | // Fluentd Forwarder 3 | // 4 | // Copyright (C) 2014 Treasure Data, Inc. 5 | // 6 | // Licensed under the Apache License, Version 2.0 (the "License"); 7 | // you may not use this file except in compliance with the License. 8 | // You may obtain a copy of the License at 9 | // 10 | // http://www.apache.org/licenses/LICENSE-2.0 11 | // 12 | // Unless required by applicable law or agreed to in writing, software 13 | // distributed under the License is distributed on an "AS IS" BASIS, 14 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | // See the License for the specific language governing permissions and 16 | // limitations under the License. 17 | // 18 | 19 | package fluentd_forwarder 20 | 21 | import ( 22 | "bytes" 23 | "compress/gzip" 24 | "encoding/hex" 25 | ioextras "github.com/moriyoshi/go-ioextras" 26 | td_client "github.com/treasure-data/td-client-go" 27 | "io/ioutil" 28 | "testing" 29 | ) 30 | 31 | func Test_CompressingBlob(t *testing.T) { 32 | for bufSize := 2; bufSize <= 32; bufSize += 1 { 33 | b := NewCompressingBlob(td_client.InMemoryBlob("testtesttesttest"), bufSize, gzip.DefaultCompression, &ioextras.MemoryRandomAccessStoreFactory{}) 34 | r, err := b.Reader() 35 | if err != nil { 36 | t.Log(err.Error()) 37 | t.FailNow() 38 | } 39 | a := [32]byte{} 40 | n, err := r.Read(a[0:4]) 41 | if err != nil { 42 | t.Log(err.Error()) 43 | t.FailNow() 44 | } 45 | t.Logf("n=%d\n", n) 46 | if n != 4 { 47 | t.Fail() 48 | } 49 | if a[0] != 0x1f || a[1] != 0x8b || a[2] != 8 || a[3] != 0 { 50 | t.Fail() 51 | } 52 | c, err := ioutil.ReadAll(r) 53 | r.Close() 54 | sum, err := b.MD5Sum() 55 | if err != nil { 56 | t.Log(err.Error()) 57 | t.FailNow() 58 | } 59 | hexSum := hex.EncodeToString(sum) 60 | if hexSum != "7e094b2f9bef89a2a889bba182e8efcf" { 61 | t.Log(hexSum) 62 | t.Fail() 63 | } 64 | s, err := b.Size() 65 | if err != nil { 66 | t.Log(err.Error()) 67 | t.FailNow() 68 | } 69 | t.Logf("size=%d", s) 70 | if s != 30 { 71 | t.Fail() 72 | } 73 | if err != nil { 74 | t.Log(err.Error()) 75 | t.FailNow() 76 | } 77 | copy(a[4:], c) 78 | t.Log(hex.EncodeToString(a[0 : 4+len(c)])) 79 | rr, err := gzip.NewReader(bytes.NewReader(a[0 : 4+len(c)])) 80 | if err != nil { 81 | t.Log(err.Error()) 82 | t.FailNow() 83 | } 84 | d, err := ioutil.ReadAll(rr) 85 | if err != nil { 86 | t.Log(err.Error()) 87 | t.FailNow() 88 | } 89 | if string(d) != "testtesttesttest" { 90 | t.Fail() 91 | } 92 | } 93 | } 94 | 95 | func Test_CompressingBlob_MD5Sum_before_reading(t *testing.T) { 96 | for bufSize := 2; bufSize <= 32; bufSize += 1 { 97 | b := NewCompressingBlob(td_client.InMemoryBlob("testtesttesttest"), bufSize, gzip.DefaultCompression, &ioextras.MemoryRandomAccessStoreFactory{}) 98 | r, err := b.Reader() 99 | if err != nil { 100 | t.Log(err.Error()) 101 | t.FailNow() 102 | } 103 | sum, err := b.MD5Sum() 104 | if err != nil { 105 | t.Log(err.Error()) 106 | t.FailNow() 107 | } 108 | hexSum := hex.EncodeToString(sum) 109 | if hexSum != "7e094b2f9bef89a2a889bba182e8efcf" { 110 | t.Log(hexSum) 111 | t.Fail() 112 | } 113 | a := [32]byte{} 114 | n, err := r.Read(a[0:4]) 115 | if err != nil { 116 | t.Log(err.Error()) 117 | t.FailNow() 118 | } 119 | t.Logf("n=%d\n", n) 120 | if n != 4 { 121 | t.Fail() 122 | } 123 | if a[0] != 0x1f || a[1] != 0x8b || a[2] != 8 || a[3] != 0 { 124 | t.Fail() 125 | } 126 | c, err := ioutil.ReadAll(r) 127 | r.Close() 128 | s, err := b.Size() 129 | if err != nil { 130 | t.Log(err.Error()) 131 | t.FailNow() 132 | } 133 | t.Logf("size=%d", s) 134 | if s != 30 { 135 | t.Fail() 136 | } 137 | if err != nil { 138 | t.Log(err.Error()) 139 | t.FailNow() 140 | } 141 | copy(a[4:], c) 142 | t.Log(hex.EncodeToString(a[0 : 4+len(c)])) 143 | rr, err := gzip.NewReader(bytes.NewReader(a[0 : 4+len(c)])) 144 | if err != nil { 145 | t.Log(err.Error()) 146 | t.FailNow() 147 | } 148 | d, err := ioutil.ReadAll(rr) 149 | if err != nil { 150 | t.Log(err.Error()) 151 | t.FailNow() 152 | } 153 | if string(d) != "testtesttesttest" { 154 | t.Fail() 155 | } 156 | } 157 | } 158 | -------------------------------------------------------------------------------- /entrypoints/build_fluentd_forwarder/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | "os/exec" 7 | "path" 8 | "path/filepath" 9 | "strings" 10 | ) 11 | 12 | func getRevision(importPath string) (string, error) { 13 | goPath := os.Getenv("GOPATH") 14 | if goPath == "" { 15 | return "", fmt.Errorf("GOPATH is not set") 16 | } 17 | repoPath := filepath.Join(goPath, "src", strings.Replace(importPath, "/", string(filepath.Separator), -1)) 18 | err := os.Chdir(repoPath) 19 | if err != nil { 20 | return "", err 21 | } 22 | cmd := exec.Command("git", "describe", "--tags", "--abbrev=0") 23 | out, err := cmd.Output() 24 | if err != nil { 25 | return "", err 26 | } 27 | return (string)(out), nil 28 | } 29 | 30 | const ImportPathBase = "github.com/fluent/fluentd-forwarder" 31 | const VersionStringVarName = "main.progVersion" 32 | 33 | func bail(message string, status int) { 34 | fmt.Fprintf(os.Stderr, "build: %s\n", message) 35 | os.Exit(status) 36 | } 37 | 38 | func main() { 39 | rev, err := getRevision(ImportPathBase) 40 | if err != nil { 41 | bail(err.Error(), 1) 42 | } 43 | for _, app := range os.Args[1:] { 44 | cmd := exec.Command("go", "get", "-u", "-ldflags", fmt.Sprintf("-X %s=%s", VersionStringVarName, rev), path.Join(ImportPathBase, "entrypoints", app)) 45 | cmd.Stdout = os.Stdout 46 | cmd.Stderr = os.Stderr 47 | err := cmd.Run() 48 | if err != nil { 49 | bail(err.Error(), 2) 50 | } 51 | } 52 | os.Exit(0) 53 | } 54 | -------------------------------------------------------------------------------- /entrypoints/dummy_td_server/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "compress/gzip" 6 | "crypto/md5" 7 | "encoding/hex" 8 | "encoding/json" 9 | "flag" 10 | "fmt" 11 | "github.com/ugorji/go/codec" 12 | "io" 13 | "io/ioutil" 14 | "net" 15 | "net/http" 16 | "os" 17 | "regexp" 18 | "time" 19 | ) 20 | 21 | type DummyServerParams struct { 22 | WriteTimeout time.Duration 23 | ReadTimeout time.Duration 24 | ListenOn string 25 | ReadThrottle int 26 | } 27 | 28 | var progName = os.Args[0] 29 | 30 | func MustParseDuration(s string) time.Duration { 31 | d, err := time.ParseDuration(s) 32 | if err != nil { 33 | panic(err) 34 | } 35 | return d 36 | } 37 | 38 | func Error(fmtStr string, args ...interface{}) { 39 | fmt.Fprint(os.Stderr, progName, ": ") 40 | fmt.Fprintf(os.Stderr, fmtStr, args...) 41 | fmt.Fprint(os.Stderr, "\n") 42 | } 43 | 44 | func ParseArgs() *DummyServerParams { 45 | readTimeout := (time.Duration)(0) 46 | writeTimeout := (time.Duration)(0) 47 | readThrottle := 0 48 | listenOn := "" 49 | 50 | flagSet := flag.NewFlagSet(progName, flag.ExitOnError) 51 | 52 | flagSet.DurationVar(&readTimeout, "read-timeout", MustParseDuration("10s"), "read timeout on wire") 53 | flagSet.DurationVar(&writeTimeout, "write-timeout", MustParseDuration("10s"), "write timeout on wire") 54 | flagSet.IntVar(&readThrottle, "read-throttle", 0, "read slottling") 55 | flagSet.StringVar(&listenOn, "listen-on", "127.0.0.1:80", "interface address and port on which the dummy server listens") 56 | flagSet.Parse(os.Args[1:]) 57 | 58 | return &DummyServerParams{ 59 | ReadTimeout: readTimeout, 60 | WriteTimeout: writeTimeout, 61 | ListenOn: listenOn, 62 | ReadThrottle: readThrottle, 63 | } 64 | } 65 | 66 | func internalServerError(resp http.ResponseWriter) { 67 | resp.WriteHeader(500) 68 | resp.Write([]byte(`{"errorMessage":"Internal Server Error"}`)) 69 | } 70 | 71 | func ReadThrottled(rdr io.Reader, l int, bps int) ([]byte, error) { 72 | _bps := int64(bps) 73 | b := make([]byte, 4096) 74 | t := time.Now() 75 | o := 0 76 | for o < l { 77 | if o+4096 >= len(b) { 78 | _b := make([]byte, len(b)+len(b)/2) 79 | copy(_b, b) 80 | b = _b 81 | } 82 | _t := time.Now() 83 | elapsed := _t.Sub(t) 84 | if elapsed > 0 { 85 | _o := int64(o) * int64(1000000000) 86 | cbps := _o / int64(elapsed) 87 | if cbps > _bps { 88 | time.Sleep(time.Duration(_o/_bps - int64(elapsed))) 89 | } 90 | } 91 | x := o + 4096 92 | if x >= len(b) { 93 | x = len(b) 94 | } 95 | n, err := rdr.Read(b[o:x]) 96 | o += n 97 | if err != nil { 98 | if err != io.EOF { 99 | return nil, err 100 | } else { 101 | break 102 | } 103 | } 104 | } 105 | b = b[0:o] 106 | return b, nil 107 | } 108 | 109 | func handleReq(params *DummyServerParams, resp http.ResponseWriter, req *http.Request, matchparams map[string]string) { 110 | resp.Header().Set("Content-Type", "application/json; charset=UTF-8") 111 | h := md5.New() 112 | format := matchparams["format"] 113 | var body []byte 114 | var err error 115 | if params.ReadThrottle > 0 { 116 | body, err = ReadThrottled(req.Body, int(req.ContentLength), params.ReadThrottle) 117 | } else { 118 | body, err = ioutil.ReadAll(req.Body) 119 | } 120 | if err != nil { 121 | internalServerError(resp) 122 | return 123 | } 124 | rdr := (io.Reader)(bytes.NewReader(body)) 125 | if format == "msgpack.gz" { 126 | rdr, err = gzip.NewReader(rdr) 127 | if err != nil { 128 | internalServerError(resp) 129 | return 130 | } 131 | } 132 | _codec := &codec.MsgpackHandle{} 133 | decoder := codec.NewDecoder(rdr, _codec) 134 | numRecords := 0 135 | for { 136 | v := map[string]interface{}{} 137 | err := decoder.Decode(&v) 138 | if err != nil { 139 | if err == io.EOF { 140 | break 141 | } else { 142 | Error("%s", err.Error()) 143 | break 144 | } 145 | } 146 | numRecords += 1 147 | } 148 | fmt.Printf("%d records received\n", numRecords) 149 | io.Copy(h, bytes.NewReader(body)) 150 | md5sum := make([]byte, 0, h.Size()) 151 | md5sum = h.Sum(md5sum) 152 | uniqueId, _ := matchparams["uniqueId"] 153 | respData := map[string]interface{}{ 154 | "unique_id": uniqueId, 155 | "database": matchparams["database"], 156 | "table": matchparams["table"], 157 | "md5_hex": hex.EncodeToString(md5sum), 158 | "elapsed_time": 0., 159 | } 160 | payload, err := json.Marshal(respData) 161 | if err != nil { 162 | internalServerError(resp) 163 | return 164 | } 165 | resp.WriteHeader(200) 166 | resp.Write(payload) 167 | } 168 | 169 | type RegexpServeMuxHandler func(http.ResponseWriter, *http.Request, map[string]string) 170 | 171 | type regexpServeMuxEntry struct { 172 | pattern *regexp.Regexp 173 | handler RegexpServeMuxHandler 174 | } 175 | 176 | type RegexpServeMux struct { 177 | entries []*regexpServeMuxEntry 178 | } 179 | 180 | func (mux *RegexpServeMux) Handle(pattern string, handler RegexpServeMuxHandler) error { 181 | rex, err := regexp.Compile(pattern) 182 | if err != nil { 183 | return err 184 | } 185 | mux.entries = append(mux.entries, ®expServeMuxEntry{ 186 | pattern: rex, 187 | handler: handler, 188 | }) 189 | return nil 190 | } 191 | 192 | func (mux *RegexpServeMux) ServeHTTP(resp http.ResponseWriter, req *http.Request) { 193 | submatches := [][]byte{} 194 | candidate := (*regexpServeMuxEntry)(nil) 195 | for _, entry := range mux.entries { 196 | submatches = entry.pattern.FindSubmatch([]byte(req.URL.Path)) 197 | if submatches != nil { 198 | candidate = entry 199 | break 200 | } 201 | } 202 | if candidate == nil { 203 | resp.WriteHeader(400) 204 | return 205 | } 206 | matchparams := map[string]string{} 207 | for i, name := range candidate.pattern.SubexpNames() { 208 | if submatches[i] != nil { 209 | // XXX: assuming URL is encoded in UTF-8 210 | matchparams[name] = string(submatches[i]) 211 | } 212 | } 213 | candidate.handler(resp, req, matchparams) 214 | } 215 | 216 | func newRegexpServeMux() *RegexpServeMux { 217 | return &RegexpServeMux{ 218 | entries: make([]*regexpServeMuxEntry, 0, 16), 219 | } 220 | } 221 | 222 | func buildMux(handle RegexpServeMuxHandler) *RegexpServeMux { 223 | mux := newRegexpServeMux() 224 | err := mux.Handle("^/v3/table/import_with_id/(?P[^/]+)/(?P[^/]+)/(?P[^/]+)/(?P[^/]+)$", handle) 225 | if err != nil { 226 | panic(err.Error()) 227 | } 228 | err = mux.Handle("^/v3/table/import/(?P[^/]+)/(?P
[^/]+)/(?P[^/]+)$", handle) 229 | if err != nil { 230 | panic(err.Error()) 231 | } 232 | return mux 233 | } 234 | 235 | func main() { 236 | params := ParseArgs() 237 | var mux = buildMux(func(resp http.ResponseWriter, req *http.Request, matchparams map[string]string) { 238 | handleReq(params, resp, req, matchparams) 239 | }) 240 | server := http.Server{ 241 | Addr: params.ListenOn, 242 | ReadTimeout: params.ReadTimeout, 243 | WriteTimeout: params.WriteTimeout, 244 | Handler: mux, 245 | } 246 | listener, err := net.Listen("tcp", params.ListenOn) 247 | if err != nil { 248 | Error("%s", err.Error()) 249 | os.Exit(1) 250 | } 251 | fmt.Printf("Dummy server listening on %s ...\n", params.ListenOn) 252 | fmt.Print("Hit CTRL-C to stop\n") 253 | server.Serve(listener) 254 | } 255 | -------------------------------------------------------------------------------- /entrypoints/fluentd_forwarder/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "crypto/x509" 5 | "flag" 6 | "fmt" 7 | fluentd_forwarder "github.com/fluent/fluentd-forwarder" 8 | strftime "github.com/jehiah/go-strftime" 9 | ioextras "github.com/moriyoshi/go-ioextras" 10 | logging "github.com/op/go-logging" 11 | gcfg "gopkg.in/gcfg.v1" 12 | "io" 13 | "io/ioutil" 14 | "log" 15 | "net/url" 16 | "os" 17 | "path/filepath" 18 | "reflect" 19 | "runtime/pprof" 20 | "strings" 21 | "time" 22 | ) 23 | 24 | type FluentdForwarderParams struct { 25 | RetryInterval time.Duration 26 | ConnectionTimeout time.Duration 27 | WriteTimeout time.Duration 28 | FlushInterval time.Duration 29 | Parallelism int 30 | JournalGroupPath string 31 | MaxJournalChunkSize int64 32 | ListenOn string 33 | OutputType string 34 | ForwardTo string 35 | LogLevel logging.Level 36 | LogFile string 37 | DatabaseName string 38 | TableName string 39 | ApiKey string 40 | Ssl bool 41 | SslCACertBundleFile string 42 | CPUProfileFile string 43 | Metadata string 44 | } 45 | 46 | type PortWorker interface { 47 | fluentd_forwarder.Port 48 | fluentd_forwarder.Worker 49 | } 50 | 51 | var progName = os.Args[0] 52 | var progVersion string 53 | 54 | func MustParseDuration(s string) time.Duration { 55 | d, err := time.ParseDuration(s) 56 | if err != nil { 57 | panic(err) 58 | } 59 | return d 60 | } 61 | 62 | func Error(fmtStr string, args ...interface{}) { 63 | fmt.Fprint(os.Stderr, progName, ": ") 64 | fmt.Fprintf(os.Stderr, fmtStr, args...) 65 | fmt.Fprint(os.Stderr, "\n") 66 | } 67 | 68 | type LogLevelValue logging.Level 69 | 70 | func (v *LogLevelValue) String() string { 71 | return logging.Level(*v).String() 72 | } 73 | 74 | func (v *LogLevelValue) Set(s string) error { 75 | _v, err := logging.LogLevel(s) 76 | *v = LogLevelValue(_v) 77 | return err 78 | } 79 | 80 | func updateFlagsByConfig(configFile string, flagSet *flag.FlagSet) error { 81 | config := struct { 82 | Fluentd_Forwarder struct { 83 | Retry_interval string `retry-interval` 84 | Conn_timeout string `conn-timeout` 85 | Write_timeout string `write-timeout` 86 | Flush_interval string `flush-interval` 87 | Listen_on string `listen-on` 88 | To string `to` 89 | Buffer_path string `buffer-path` 90 | Buffer_chunk_limit string `buffer-chunk-limit` 91 | Log_level string `log-level` 92 | Ca_certs string `ca-certs` 93 | Cpuprofile string `cpuprofile` 94 | Log_file string `log-file` 95 | } 96 | }{} 97 | err := gcfg.ReadFileInto(&config, configFile) 98 | if err != nil { 99 | return err 100 | } 101 | r := reflect.ValueOf(config.Fluentd_Forwarder) 102 | rt := r.Type() 103 | for i, l := 0, rt.NumField(); i < l; i += 1 { 104 | f := rt.Field(i) 105 | fv := r.Field(i) 106 | v := fv.String() 107 | if v != "" { 108 | err := flagSet.Set(string(f.Tag), v) 109 | if err != nil { 110 | return err 111 | } 112 | } 113 | } 114 | return nil 115 | } 116 | 117 | func ParseArgs() *FluentdForwarderParams { 118 | configFile := "" 119 | retryInterval := (time.Duration)(0) 120 | connectionTimeout := (time.Duration)(0) 121 | writeTimeout := (time.Duration)(0) 122 | flushInterval := (time.Duration)(0) 123 | parallelism := 0 124 | listenOn := "" 125 | forwardTo := "" 126 | journalGroupPath := "" 127 | maxJournalChunkSize := int64(16777216) 128 | logLevel := LogLevelValue(logging.INFO) 129 | sslCACertBundleFile := "" 130 | cpuProfileFile := "" 131 | logFile := "" 132 | metadata := "" 133 | 134 | flagSet := flag.NewFlagSet(progName, flag.ExitOnError) 135 | 136 | flagSet.StringVar(&configFile, "config", "", "configuration file") 137 | flagSet.DurationVar(&retryInterval, "retry-interval", 0, "retry interval in which connection is tried against the remote agent") 138 | flagSet.DurationVar(&connectionTimeout, "conn-timeout", MustParseDuration("10s"), "connection timeout") 139 | flagSet.DurationVar(&writeTimeout, "write-timeout", MustParseDuration("10s"), "write timeout on wire") 140 | flagSet.DurationVar(&flushInterval, "flush-interval", MustParseDuration("5s"), "flush interval in which the events are forwareded to the remote agent") 141 | flagSet.IntVar(¶llelism, "parallelism", 1, "Number of chunks to submit at once (for td output)") 142 | flagSet.StringVar(&listenOn, "listen-on", "127.0.0.1:24224", "interface address and port on which the forwarder listens") 143 | flagSet.StringVar(&forwardTo, "to", "fluent://127.0.0.1:24225", "host and port to which the events are forwarded") 144 | flagSet.StringVar(&journalGroupPath, "buffer-path", "*", "directory / path on which buffer files are created. * may be used within the path to indicate the prefix or suffix like var/pre*suf") 145 | flagSet.Int64Var(&maxJournalChunkSize, "buffer-chunk-limit", 16777216, "Maximum size of a buffer chunk") 146 | flagSet.Var(&logLevel, "log-level", "log level (defaults to INFO)") 147 | flagSet.StringVar(&sslCACertBundleFile, "ca-certs", "", "path to SSL CA certificate bundle file") 148 | flagSet.StringVar(&cpuProfileFile, "cpuprofile", "", "write CPU profile to file") 149 | flagSet.StringVar(&logFile, "log-file", "", "path of the log file. log will be written to stderr if unspecified") 150 | flagSet.StringVar(&metadata, "metadata", "", "set addtional data into record") 151 | flagSet.Parse(os.Args[1:]) 152 | 153 | if configFile != "" { 154 | err := updateFlagsByConfig(configFile, flagSet) 155 | if err != nil { 156 | Error("%s", err.Error()) 157 | os.Exit(1) 158 | } 159 | } 160 | 161 | ssl := false 162 | outputType := "" 163 | databaseName := "*" 164 | tableName := "*" 165 | apiKey := "" 166 | 167 | if strings.Contains(forwardTo, "//") { 168 | u, err := url.Parse(forwardTo) 169 | if err != nil { 170 | Error("%s", err.Error()) 171 | os.Exit(1) 172 | } 173 | switch u.Scheme { 174 | case "fluent", "fluentd": 175 | outputType = "fluent" 176 | forwardTo = u.Host 177 | case "td+http", "td+https": 178 | outputType = "td" 179 | forwardTo = u.Host 180 | if u.User != nil { 181 | apiKey = u.User.Username() 182 | } 183 | if u.Scheme == "td+https" { 184 | ssl = true 185 | } 186 | p := strings.Split(u.Path, "/") 187 | if len(p) > 1 { 188 | databaseName = p[1] 189 | } 190 | if len(p) > 2 { 191 | tableName = p[2] 192 | } 193 | } 194 | } else { 195 | outputType = "fluent" 196 | } 197 | if outputType == "" { 198 | Error("Invalid output specifier") 199 | os.Exit(1) 200 | } else if outputType == "fluent" { 201 | if !strings.ContainsRune(forwardTo, ':') { 202 | forwardTo += ":24224" 203 | } 204 | } 205 | return &FluentdForwarderParams{ 206 | RetryInterval: retryInterval, 207 | ConnectionTimeout: connectionTimeout, 208 | WriteTimeout: writeTimeout, 209 | FlushInterval: flushInterval, 210 | Parallelism: parallelism, 211 | ListenOn: listenOn, 212 | OutputType: outputType, 213 | ForwardTo: forwardTo, 214 | Ssl: ssl, 215 | DatabaseName: databaseName, 216 | TableName: tableName, 217 | ApiKey: apiKey, 218 | JournalGroupPath: journalGroupPath, 219 | MaxJournalChunkSize: maxJournalChunkSize, 220 | LogLevel: logging.Level(logLevel), 221 | LogFile: logFile, 222 | SslCACertBundleFile: sslCACertBundleFile, 223 | CPUProfileFile: cpuProfileFile, 224 | Metadata: metadata, 225 | } 226 | } 227 | 228 | func ValidateParams(params *FluentdForwarderParams) bool { 229 | if params.RetryInterval < 0 { 230 | Error("Retry interval may not be negative") 231 | return false 232 | } 233 | if params.RetryInterval > 0 && params.RetryInterval < 100000000 { 234 | Error("Retry interval must be greater than or equal to 100ms") 235 | return false 236 | } 237 | if params.FlushInterval < 100000000 { 238 | Error("Flush interval must be greater than or equal to 100ms") 239 | return false 240 | } 241 | if params.FlushInterval < 100000000 { 242 | Error("Flush interval must be greater than or equal to 100ms") 243 | return false 244 | } 245 | switch params.OutputType { 246 | case "fluent": 247 | if params.RetryInterval == 0 { 248 | params.RetryInterval = MustParseDuration("5s") 249 | } 250 | if params.RetryInterval > params.FlushInterval { 251 | Error("Retry interval may not be greater than flush interval") 252 | return false 253 | } 254 | case "td": 255 | if params.RetryInterval != 0 { 256 | Error("Retry interval will be ignored") 257 | return false 258 | } 259 | } 260 | return true 261 | } 262 | 263 | func main() { 264 | params := ParseArgs() 265 | if !ValidateParams(params) { 266 | os.Exit(1) 267 | } 268 | logWriter := (io.Writer)(nil) 269 | if params.LogFile != "" { 270 | logWriter = ioextras.NewStaticRotatingWriter( 271 | func(_ interface{}) (string, error) { 272 | path := strftime.Format(params.LogFile, time.Now()) 273 | return path, nil 274 | }, 275 | func(path string, _ interface{}) (io.Writer, error) { 276 | dir, _ := filepath.Split(path) 277 | err := os.MkdirAll(dir, os.FileMode(0777)) 278 | if err != nil { 279 | return nil, err 280 | } 281 | return os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, os.FileMode(0666)) 282 | }, 283 | nil, 284 | ) 285 | } else { 286 | logWriter = os.Stderr 287 | } 288 | logBackend := logging.NewLogBackend(logWriter, "[fluentd-forwarder] ", log.Ldate|log.Ltime|log.Lmicroseconds) 289 | logging.SetBackend(logBackend) 290 | logger := logging.MustGetLogger("fluentd-forwarder") 291 | logging.SetLevel(params.LogLevel, "fluentd-forwarder") 292 | if progVersion != "" { 293 | logger.Infof("Version %s starting...", progVersion) 294 | } 295 | 296 | workerSet := fluentd_forwarder.NewWorkerSet() 297 | 298 | if params.CPUProfileFile != "" { 299 | f, err := os.Create(params.CPUProfileFile) 300 | if err != nil { 301 | Error(err.Error()) 302 | os.Exit(1) 303 | } 304 | pprof.StartCPUProfile(f) 305 | defer pprof.StopCPUProfile() 306 | } 307 | 308 | output := (PortWorker)(nil) 309 | err := (error)(nil) 310 | switch params.OutputType { 311 | case "fluent": 312 | output, err = fluentd_forwarder.NewForwardOutput( 313 | logger, 314 | params.ForwardTo, 315 | params.RetryInterval, 316 | params.ConnectionTimeout, 317 | params.WriteTimeout, 318 | params.FlushInterval, 319 | params.JournalGroupPath, 320 | params.MaxJournalChunkSize, 321 | params.Metadata, 322 | ) 323 | case "td": 324 | rootCAs := (*x509.CertPool)(nil) 325 | if params.SslCACertBundleFile != "" { 326 | b, err := ioutil.ReadFile(params.SslCACertBundleFile) 327 | if err != nil { 328 | Error("Failed to read CA bundle file: %s", err.Error()) 329 | os.Exit(1) 330 | } 331 | rootCAs = x509.NewCertPool() 332 | if !rootCAs.AppendCertsFromPEM(b) { 333 | Error("No valid certificate found in %s", params.SslCACertBundleFile) 334 | os.Exit(1) 335 | } 336 | } 337 | output, err = fluentd_forwarder.NewTDOutput( 338 | logger, 339 | params.ForwardTo, 340 | params.ConnectionTimeout, 341 | params.WriteTimeout, 342 | params.FlushInterval, 343 | params.Parallelism, 344 | params.JournalGroupPath, 345 | params.MaxJournalChunkSize, 346 | params.ApiKey, 347 | params.DatabaseName, 348 | params.TableName, 349 | "", 350 | params.Ssl, 351 | rootCAs, 352 | "", // TODO:http-proxy 353 | params.Metadata, 354 | ) 355 | } 356 | if err != nil { 357 | Error("%s", err.Error()) 358 | return 359 | } 360 | workerSet.Add(output) 361 | input, err := fluentd_forwarder.NewForwardInput(logger, params.ListenOn, output) 362 | if err != nil { 363 | Error(err.Error()) 364 | return 365 | } 366 | workerSet.Add(input) 367 | 368 | signalHandler := NewSignalHandler(workerSet) 369 | input.Start() 370 | output.Start() 371 | signalHandler.Start() 372 | 373 | for _, worker := range workerSet.Slice() { 374 | worker.WaitForShutdown() 375 | } 376 | logger.Notice("Shutting down...") 377 | } 378 | -------------------------------------------------------------------------------- /entrypoints/fluentd_forwarder/signal.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | fluentd_forwarder "github.com/fluent/fluentd-forwarder" 5 | "os" 6 | "os/signal" 7 | ) 8 | 9 | type SignalHandler struct { 10 | Workers *fluentd_forwarder.WorkerSet 11 | signalChan chan os.Signal 12 | } 13 | 14 | func (handler *SignalHandler) Start() { 15 | signal.Notify(handler.signalChan, os.Kill, os.Interrupt) 16 | go func() { 17 | <-handler.signalChan 18 | for _, worker := range handler.Workers.Slice() { 19 | worker.Stop() 20 | } 21 | }() 22 | } 23 | 24 | func NewSignalHandler(workerSet *fluentd_forwarder.WorkerSet) *SignalHandler { 25 | return &SignalHandler{ 26 | workerSet, 27 | make(chan os.Signal, 1), 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /errors.go: -------------------------------------------------------------------------------- 1 | // 2 | // Fluentd Forwarder 3 | // 4 | // Copyright (C) 2014 Treasure Data, Inc. 5 | // 6 | // Licensed under the Apache License, Version 2.0 (the "License"); 7 | // you may not use this file except in compliance with the License. 8 | // You may obtain a copy of the License at 9 | // 10 | // http://www.apache.org/licenses/LICENSE-2.0 11 | // 12 | // Unless required by applicable law or agreed to in writing, software 13 | // distributed under the License is distributed on an "AS IS" BASIS, 14 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | // See the License for the specific language governing permissions and 16 | // limitations under the License. 17 | // 18 | 19 | package fluentd_forwarder 20 | 21 | import ( 22 | "fmt" 23 | ) 24 | 25 | type Errors []error 26 | 27 | func (e Errors) Error() string { 28 | buf := []byte("Failure due to one or more errors: ") 29 | for i, e_ := range e { 30 | if i > 0 { 31 | buf = append(buf, " / "...) 32 | } 33 | buf = append(buf, fmt.Sprintf("%d: %s", i+1, e_.Error())...) 34 | } 35 | return string(buf) 36 | } 37 | -------------------------------------------------------------------------------- /file_journal.go: -------------------------------------------------------------------------------- 1 | // 2 | // Copyright (C) 2014 Moriyoshi Koizumi 3 | // Copyright (C) 2014 Treasure Data, Inc. 4 | // 5 | // Permission is hereby granted, free of charge, to any person obtaining a copy 6 | // of this software and associated documentation files (the "Software"), to deal 7 | // in the Software without restriction, including without limitation the rights 8 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | // copies of the Software, and to permit persons to whom the Software is 10 | // furnished to do so, subject to the following conditions: 11 | // 12 | // The above copyright notice and this permission notice shall be included in 13 | // all copies or substantial portions of the Software. 14 | // 15 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 21 | // THE SOFTWARE. 22 | // 23 | 24 | package fluentd_forwarder 25 | 26 | import ( 27 | "crypto/md5" 28 | "encoding/hex" 29 | "errors" 30 | "fmt" 31 | logging "github.com/op/go-logging" 32 | "io" 33 | "math/rand" 34 | "os" 35 | "path/filepath" 36 | "strings" 37 | "sync" 38 | "sync/atomic" 39 | "time" 40 | "unsafe" 41 | ) 42 | 43 | type FileJournalChunkDequeueHead struct { 44 | next *FileJournalChunk 45 | prev *FileJournalChunk 46 | } 47 | 48 | type FileJournalChunkDequeue struct { 49 | first *FileJournalChunk 50 | last *FileJournalChunk 51 | count int 52 | mtx sync.Mutex 53 | } 54 | 55 | type FileJournalChunk struct { 56 | Size int64 // This variable must be on 64-bit alignment. Otherwise atomic.AddInt64 will cause a crash on ARM and x86-32 57 | head FileJournalChunkDequeueHead 58 | container *FileJournalChunkDequeue 59 | Path string 60 | Type JournalFileType 61 | TSuffix string 62 | Timestamp int64 63 | UniqueId []byte 64 | refcount int32 65 | mtx sync.Mutex 66 | } 67 | 68 | type FileJournal struct { 69 | group *FileJournalGroup 70 | key string 71 | chunks FileJournalChunkDequeue 72 | writer io.WriteCloser 73 | newChunkListeners map[JournalChunkListener]JournalChunkListener 74 | flushListeners map[JournalChunkListener]JournalChunkListener 75 | mtx sync.Mutex 76 | } 77 | 78 | type FileJournalGroup struct { 79 | factory *FileJournalGroupFactory 80 | worker Worker 81 | timeGetter func() time.Time 82 | logger *logging.Logger 83 | rand *rand.Rand 84 | fileMode os.FileMode 85 | maxSize int64 86 | pathPrefix string 87 | pathSuffix string 88 | journals map[string]*FileJournal 89 | mtx sync.Mutex 90 | } 91 | 92 | type FileJournalGroupFactory struct { 93 | logger *logging.Logger 94 | paths map[string]*FileJournalGroup 95 | randSource rand.Source 96 | timeGetter func() time.Time 97 | defaultPathSuffix string 98 | defaultFileMode os.FileMode 99 | maxSize int64 100 | } 101 | 102 | type FileJournalChunkWrapper struct { 103 | journal *FileJournal 104 | chunk *FileJournalChunk 105 | } 106 | 107 | func (wrapper *FileJournalChunkWrapper) Path() (string, error) { 108 | chunk := (*FileJournalChunk)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&wrapper.chunk)))) 109 | if chunk == nil { 110 | return "", errors.New("already disposed") 111 | } 112 | return chunk.getPath(), nil 113 | } 114 | 115 | func (wrapper *FileJournalChunkWrapper) Id() string { 116 | chunk := (*FileJournalChunk)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&wrapper.chunk)))) 117 | if chunk == nil { 118 | return "" 119 | } 120 | return hex.EncodeToString(chunk.UniqueId) 121 | } 122 | 123 | func (wrapper *FileJournalChunkWrapper) String() string { 124 | retval, err := wrapper.Path() 125 | if err != nil { 126 | return err.Error() 127 | } else { 128 | return retval 129 | } 130 | } 131 | 132 | func (wrapper *FileJournalChunkWrapper) Size() (int64, error) { 133 | chunk := (*FileJournalChunk)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&wrapper.chunk)))) 134 | if chunk == nil { 135 | return -1, errors.New("already disposed") 136 | } 137 | return chunk.getSize(), nil 138 | } 139 | 140 | func (wrapper *FileJournalChunkWrapper) Reader() (io.ReadCloser, error) { 141 | chunk := (*FileJournalChunk)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&wrapper.chunk)))) 142 | if chunk == nil { 143 | return nil, errors.New("already disposed") 144 | } 145 | return chunk.getReader() 146 | } 147 | 148 | func (wrapper *FileJournalChunkWrapper) MD5Sum() ([]byte, error) { 149 | chunk := (*FileJournalChunk)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&wrapper.chunk)))) 150 | if chunk == nil { 151 | return nil, errors.New("already disposed") 152 | } 153 | return chunk.md5Sum() 154 | } 155 | 156 | func (wrapper *FileJournalChunkWrapper) NextChunk() JournalChunk { 157 | chunk := (*FileJournalChunk)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&wrapper.chunk)))) 158 | if chunk == nil { 159 | return nil 160 | } 161 | nextChunk := chunk.getNextChunk(wrapper.journal) 162 | if nextChunk != nil { 163 | return wrapper.journal.newChunkWrapper(nextChunk) 164 | } else { 165 | return nil 166 | } 167 | } 168 | 169 | func (wrapper *FileJournalChunkWrapper) Dispose() error { 170 | chunk := (*FileJournalChunk)(atomic.SwapPointer((*unsafe.Pointer)(unsafe.Pointer(&wrapper.chunk)), nil)) 171 | if chunk == nil { 172 | return errors.New("already disposed") 173 | } 174 | return wrapper.journal.deleteRef((*FileJournalChunk)(chunk)) 175 | } 176 | 177 | func (wrapper *FileJournalChunkWrapper) Dup() JournalChunk { 178 | chunk := (*FileJournalChunk)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&wrapper.chunk)))) 179 | if chunk == nil { 180 | return nil 181 | } 182 | return wrapper.journal.newChunkWrapper(chunk) 183 | } 184 | 185 | func (journal *FileJournal) newChunkWrapper(chunk *FileJournalChunk) *FileJournalChunkWrapper { 186 | journal.addRef(chunk) 187 | return &FileJournalChunkWrapper{journal, chunk} 188 | } 189 | 190 | func (journal *FileJournal) addRef(chunk *FileJournalChunk) int32 { 191 | return atomic.AddInt32(&chunk.refcount, 1) 192 | } 193 | 194 | func (journal *FileJournal) deleteRef(chunk *FileJournalChunk) error { 195 | refcount := atomic.AddInt32(&chunk.refcount, -1) 196 | if refcount == 0 { 197 | chunk.mtx.Lock() 198 | defer chunk.mtx.Unlock() 199 | container := (*FileJournalChunkDequeue)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&chunk.container)))) 200 | container.mtx.Lock() 201 | defer container.mtx.Unlock() 202 | err := os.Remove(chunk.Path) 203 | if err != nil { 204 | // undo the change 205 | atomic.AddInt32(&chunk.refcount, 1) 206 | return err 207 | } 208 | { 209 | prevChunk := chunk.head.prev 210 | nextChunk := chunk.head.next 211 | if prevChunk != nil { 212 | prevChunk.head.next = nextChunk 213 | } else if container.first == chunk { 214 | container.first = nextChunk 215 | } 216 | if nextChunk != nil { 217 | nextChunk.head.prev = prevChunk 218 | } else if container.last == chunk { 219 | container.last = prevChunk 220 | } 221 | chunk.head.prev = nil 222 | chunk.head.next = nil 223 | container.count -= 1 224 | } 225 | return nil 226 | } else if refcount < 0 { 227 | // should never happen 228 | panic(fmt.Sprintf("something went wrong! chunk=%s, chunks.count=%d", chunk.Path, chunk.container.count)) 229 | } 230 | return nil 231 | } 232 | 233 | func (chunk *FileJournalChunk) getReader() (io.ReadCloser, error) { 234 | chunk.mtx.Lock() 235 | defer chunk.mtx.Unlock() 236 | rdr, err := os.OpenFile(chunk.Path, os.O_RDONLY, 0) 237 | if err != nil { 238 | return nil, err 239 | } 240 | return rdr, err 241 | } 242 | 243 | func (chunk *FileJournalChunk) getPath() string { 244 | chunk.mtx.Lock() 245 | defer chunk.mtx.Unlock() 246 | return chunk.Path 247 | } 248 | 249 | func (chunk *FileJournalChunk) getSize() int64 { 250 | return atomic.LoadInt64(&chunk.Size) 251 | } 252 | 253 | func (chunk *FileJournalChunk) md5Sum() ([]byte, error) { 254 | chunk.mtx.Lock() 255 | defer chunk.mtx.Unlock() 256 | h := md5.New() 257 | rdr, err := os.OpenFile(chunk.Path, os.O_RDONLY, 0) 258 | if err != nil { 259 | return nil, err 260 | } 261 | defer rdr.Close() 262 | _, err = io.Copy(h, rdr) 263 | if err != nil { 264 | return nil, err 265 | } 266 | retval := make([]byte, 0, h.Size()) 267 | return h.Sum(retval), err 268 | } 269 | 270 | func (chunk *FileJournalChunk) getNextChunk(journal *FileJournal) *FileJournalChunk { 271 | chunk.mtx.Lock() 272 | defer chunk.mtx.Unlock() 273 | container := (*FileJournalChunkDequeue)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&chunk.container)))) 274 | container.mtx.Lock() 275 | defer container.mtx.Unlock() 276 | return chunk.head.prev 277 | } 278 | 279 | func (journal *FileJournal) Key() string { 280 | return journal.key 281 | } 282 | 283 | func (journal *FileJournal) notifyFlushListeners(chunk *FileJournalChunk) { 284 | // lock for listener container must be acquired by caller 285 | for _, listener := range journal.flushListeners { 286 | err := listener.ChunkFlushed(journal.newChunkWrapper(chunk)) 287 | if err != nil { 288 | journal.group.logger.Errorf("error occurred during notifying flush event: %s", err.Error()) 289 | } 290 | } 291 | } 292 | 293 | func (journal *FileJournal) notifyNewChunkListeners(chunk *FileJournalChunk) { 294 | // lock for listener container must be acquired by caller 295 | for _, listener := range journal.newChunkListeners { 296 | err := listener.NewChunkCreated(journal.newChunkWrapper(chunk)) 297 | if err != nil { 298 | journal.group.logger.Errorf("error occurred during notifying flush event: %s", err.Error()) 299 | } 300 | } 301 | } 302 | 303 | func (journal *FileJournal) finalizeChunk(chunk *FileJournalChunk) error { 304 | group := journal.group 305 | variablePortion := BuildJournalPathWithTSuffix( 306 | journal.key, 307 | Rest, 308 | chunk.TSuffix, 309 | ) 310 | newPath := group.pathPrefix + variablePortion + group.pathSuffix 311 | err := func() error { 312 | chunk.mtx.Lock() 313 | defer chunk.mtx.Unlock() 314 | err := os.Rename(chunk.Path, newPath) 315 | if err != nil { 316 | return err 317 | } 318 | chunk.Type = Rest 319 | chunk.Path = newPath 320 | return nil 321 | }() 322 | if err != nil { 323 | return err 324 | } 325 | journal.notifyFlushListeners(chunk) 326 | return nil 327 | } 328 | 329 | func (journal *FileJournal) Flush(visitor func(JournalChunk) interface{}) error { 330 | err := func() error { 331 | journal.mtx.Lock() 332 | defer journal.mtx.Unlock() 333 | // this is safe the journal lock prevents any new chunk from being added to the head 334 | head := journal.chunks.first 335 | if head != nil { 336 | _, err := journal.newChunk() 337 | if err != nil { 338 | return err 339 | } 340 | } 341 | return nil 342 | }() 343 | if err != nil { 344 | return err 345 | } 346 | dequeue := func() *FileJournalChunkDequeue { 347 | journal.chunks.mtx.Lock() 348 | defer journal.chunks.mtx.Unlock() 349 | // detach all the following chunks 350 | firstChunk := journal.chunks.first 351 | lastChunk := journal.chunks.last 352 | if firstChunk != lastChunk { 353 | nextOfFirstChunk := firstChunk.head.next 354 | dequeue := &FileJournalChunkDequeue{ 355 | first: nextOfFirstChunk, 356 | last: lastChunk, 357 | count: journal.chunks.count - 1, 358 | } 359 | firstChunk.head.next = nil 360 | nextOfFirstChunk.head.prev = nil 361 | journal.chunks.last = journal.chunks.first 362 | journal.chunks.count = 1 363 | for chunk := nextOfFirstChunk; chunk != nil; chunk = chunk.head.next { 364 | atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&chunk.container)), unsafe.Pointer(dequeue)) 365 | } 366 | return dequeue 367 | } 368 | return nil 369 | }() 370 | if dequeue != nil { 371 | journal.group.logger.Debugf("chunks to flush: %d", dequeue.count) 372 | type pair struct { 373 | chunk *FileJournalChunk 374 | futureErr <-chan error 375 | } 376 | if visitor != nil { 377 | pairs := make([]pair, 0, dequeue.count) 378 | prevChunk := (*FileJournalChunk)(nil) 379 | for chunk := dequeue.last; chunk != nil; chunk = prevChunk { 380 | prevChunk = chunk.head.prev 381 | errOrFuture := visitor(journal.newChunkWrapper(chunk)) 382 | if errOrFuture != nil { 383 | // either synchronous or asynchronous 384 | var ok bool 385 | var futureErr <-chan error 386 | err, ok = errOrFuture.(error) 387 | if ok { 388 | // synchronous 389 | _futureErr := make(chan error, 1) 390 | _futureErr <- err 391 | futureErr = _futureErr 392 | } else { 393 | // asynchrnous 394 | futureErr, ok = errOrFuture.(<-chan error) 395 | if !ok { 396 | panic("visitor returned something that is neither an error nor a channel") 397 | } 398 | } 399 | pairs = append(pairs, pair{chunk, futureErr}) 400 | } else { 401 | // synchronous mode 402 | err = journal.deleteRef(chunk) 403 | if err != nil { 404 | futureErr := make(chan error, 1) 405 | futureErr <- err 406 | pairs = append(pairs, pair{chunk, futureErr}) 407 | } 408 | 409 | } 410 | } 411 | errors := make(Errors, 0, len(pairs)) 412 | for _, p := range pairs { 413 | err := <-p.futureErr 414 | if err != nil { 415 | errors = append(errors, err) 416 | } else { 417 | err = journal.deleteRef(p.chunk) 418 | if err != nil { 419 | errors = append(errors, err) 420 | } 421 | } 422 | } 423 | journal.group.logger.Debugf("errors=%d, chunks=%d", len(errors), dequeue.count) 424 | if len(errors) > 0 { 425 | err = errors 426 | } 427 | } else { 428 | prevChunk := (*FileJournalChunk)(nil) 429 | for chunk := dequeue.last; chunk != nil; chunk = prevChunk { 430 | prevChunk = chunk.head.prev 431 | journal.deleteRef(chunk) 432 | } 433 | } 434 | func() { 435 | // re-attach chunks 436 | journal.chunks.mtx.Lock() 437 | defer journal.chunks.mtx.Unlock() 438 | dequeue.mtx.Lock() 439 | defer dequeue.mtx.Unlock() 440 | for chunk := dequeue.last; chunk != nil; chunk = chunk.head.prev { 441 | atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&chunk.container)), unsafe.Pointer(&journal.chunks)) 442 | } 443 | if journal.chunks.last != nil { 444 | journal.chunks.last.head.next = dequeue.first 445 | } else { 446 | journal.chunks.first = dequeue.first 447 | } 448 | if dequeue.first != nil { 449 | dequeue.first.head.prev = journal.chunks.last 450 | journal.chunks.last = dequeue.last 451 | } 452 | journal.chunks.count += dequeue.count 453 | }() 454 | } 455 | return err 456 | } 457 | 458 | func (journal *FileJournal) newChunk() (*FileJournalChunk, error) { 459 | group := journal.group 460 | info := BuildJournalPath( 461 | journal.key, 462 | Head, 463 | group.timeGetter(), 464 | group.rand.Int63n(0xfff), 465 | ) 466 | chunk := &FileJournalChunk{ 467 | head: FileJournalChunkDequeueHead{journal.chunks.first, nil}, 468 | container: &journal.chunks, 469 | Path: (group.pathPrefix + info.VariablePortion + group.pathSuffix), 470 | Type: info.Type, 471 | TSuffix: info.TSuffix, 472 | UniqueId: info.UniqueId, 473 | refcount: 1, 474 | } 475 | file, err := os.OpenFile(chunk.Path, os.O_WRONLY|os.O_APPEND|os.O_CREATE|os.O_EXCL, journal.group.fileMode) 476 | if err != nil { 477 | return nil, err 478 | } 479 | if journal.writer != nil { 480 | err := journal.writer.Close() 481 | if err != nil { 482 | return nil, err 483 | } 484 | journal.writer = nil 485 | } 486 | 487 | oldHead := (*FileJournalChunk)(nil) 488 | { 489 | journal.chunks.mtx.Lock() 490 | oldHead = journal.chunks.first 491 | if oldHead != nil { 492 | oldHead.head.prev = chunk 493 | } else { 494 | journal.chunks.last = chunk 495 | } 496 | chunk.head.next = oldHead 497 | journal.chunks.first = chunk 498 | journal.chunks.count += 1 499 | journal.chunks.mtx.Unlock() 500 | } 501 | chunk.refcount += 1 // for writer 502 | 503 | if oldHead != nil { 504 | err := journal.finalizeChunk(oldHead) 505 | if err != nil { 506 | file.Close() 507 | os.Remove(chunk.Path) 508 | return nil, err 509 | } 510 | err = journal.deleteRef(oldHead) // writer-holding ref 511 | if err != nil { 512 | file.Close() 513 | os.Remove(chunk.Path) 514 | return nil, err 515 | } 516 | } 517 | 518 | journal.writer = file 519 | journal.chunks.first.Size = 0 520 | journal.notifyNewChunkListeners(chunk) 521 | return chunk, nil 522 | } 523 | 524 | func (journal *FileJournal) AddFlushListener(listener JournalChunkListener) { 525 | journal.mtx.Lock() 526 | defer journal.mtx.Unlock() 527 | journal.flushListeners[listener] = listener 528 | } 529 | 530 | func (journal *FileJournal) AddNewChunkListener(listener JournalChunkListener) { 531 | journal.mtx.Lock() 532 | defer journal.mtx.Unlock() 533 | journal.newChunkListeners[listener] = listener 534 | } 535 | 536 | func (journal *FileJournal) Write(data []byte) error { 537 | journal.mtx.Lock() 538 | defer journal.mtx.Unlock() 539 | 540 | newChunkNeeded := false 541 | { 542 | journal.chunks.mtx.Lock() 543 | newChunkNeeded = journal.writer == nil || journal.chunks.first == nil || journal.group.maxSize-journal.chunks.first.Size < int64(len(data)) 544 | journal.chunks.mtx.Unlock() 545 | } 546 | if newChunkNeeded { 547 | _, err := journal.newChunk() 548 | if err != nil { 549 | return err 550 | } 551 | } 552 | if journal.writer == nil { 553 | return errors.New("journal has been disposed?") 554 | } 555 | n, err := journal.writer.Write(data) 556 | if err != nil { 557 | return err 558 | } 559 | if n != len(data) { 560 | return errors.New("not all data could be written") 561 | } 562 | atomic.AddInt64(&journal.chunks.first.Size, int64(n)) 563 | return nil 564 | } 565 | 566 | func (journal *FileJournal) TailChunk() JournalChunk { 567 | retval := (*FileJournalChunkWrapper)(nil) 568 | { 569 | journal.chunks.mtx.Lock() 570 | if journal.chunks.last != nil { 571 | retval = journal.newChunkWrapper(journal.chunks.last) 572 | } 573 | journal.chunks.mtx.Unlock() 574 | } 575 | return retval 576 | } 577 | 578 | func (journal *FileJournal) Dispose() error { 579 | journal.mtx.Lock() 580 | defer journal.mtx.Unlock() 581 | if journal.writer != nil { 582 | err := journal.writer.Close() 583 | if err != nil { 584 | return err 585 | } 586 | journal.writer = nil 587 | if journal.chunks.first != nil { 588 | err := journal.deleteRef(journal.chunks.first) 589 | if err != nil { 590 | return err 591 | } 592 | } 593 | } 594 | return nil 595 | } 596 | 597 | func (journalGroup *FileJournalGroup) Dispose() error { 598 | for _, journal := range journalGroup.journals { 599 | journal.Dispose() 600 | } 601 | return nil 602 | } 603 | 604 | func (journalGroup *FileJournalGroup) GetFileJournal(key string) *FileJournal { 605 | journalGroup.mtx.Lock() 606 | defer journalGroup.mtx.Unlock() 607 | 608 | journal, ok := journalGroup.journals[key] 609 | if ok { 610 | return journal 611 | } 612 | journal = &FileJournal{ 613 | group: journalGroup, 614 | key: key, 615 | chunks: FileJournalChunkDequeue{nil, nil, 0, sync.Mutex{}}, 616 | writer: nil, 617 | newChunkListeners: make(map[JournalChunkListener]JournalChunkListener), 618 | flushListeners: make(map[JournalChunkListener]JournalChunkListener), 619 | } 620 | journalGroup.journals[key] = journal 621 | return journal 622 | } 623 | 624 | func (journalGroup *FileJournalGroup) GetJournal(key string) Journal { 625 | return journalGroup.GetFileJournal(key) 626 | } 627 | 628 | func (journalGroup *FileJournalGroup) GetJournalKeys() []string { 629 | journalGroup.mtx.Lock() 630 | defer journalGroup.mtx.Unlock() 631 | 632 | retval := make([]string, len(journalGroup.journals)) 633 | i := 0 634 | for k := range journalGroup.journals { 635 | retval[i] = k 636 | i += 1 637 | } 638 | return retval 639 | } 640 | 641 | // http://stackoverflow.com/questions/1525117/whats-the-fastest-algorithm-for-sorting-a-linked-list 642 | // http://www.chiark.greenend.org.uk/~sgtatham/algorithms/listsort.html 643 | func sortChunksByTimestamp(chunks *FileJournalChunkDequeue) { 644 | k := 1 645 | lhs := chunks.first 646 | if lhs == nil { 647 | return 648 | } 649 | for { 650 | result := FileJournalChunkDequeue{nil, nil, chunks.count, sync.Mutex{}} 651 | first := true 652 | for { 653 | picked := (*FileJournalChunk)(nil) 654 | lhsSize := 0 655 | rhsSize := k 656 | rhs := lhs 657 | i := k 658 | for i > 0 && rhs.head.next != nil { 659 | i -= 1 660 | rhs = rhs.head.next 661 | } 662 | lhsSize = k - i 663 | for { 664 | if lhsSize != 0 { 665 | if rhsSize != 0 && rhs != nil && lhs.Timestamp < rhs.Timestamp { 666 | picked = rhs 667 | rhs = rhs.head.next 668 | rhsSize -= 1 669 | } else { 670 | picked = lhs 671 | lhs = lhs.head.next 672 | lhsSize -= 1 673 | } 674 | } else { 675 | if rhsSize != 0 && rhs != nil { 676 | picked = rhs 677 | rhs = rhs.head.next 678 | rhsSize -= 1 679 | } else { 680 | break 681 | } 682 | } 683 | if picked.head.prev != nil { 684 | picked.head.prev.head.next = picked.head.next 685 | } 686 | if picked.head.next != nil { 687 | picked.head.next.head.prev = picked.head.prev 688 | } 689 | if result.last == nil { 690 | result.first = picked 691 | } else { 692 | result.last.head.next = picked 693 | } 694 | picked.head.prev = result.last 695 | picked.head.next = nil 696 | result.last = picked 697 | } 698 | lhs = rhs 699 | if lhs == nil { 700 | break 701 | } 702 | first = false 703 | } 704 | if first { 705 | *chunks = result 706 | break 707 | } 708 | k *= 2 709 | lhs = result.first 710 | } 711 | } 712 | 713 | func validateChunks(chunks *FileJournalChunkDequeue) error { 714 | chunkHead := (*FileJournalChunk)(nil) 715 | for chunk := chunks.first; chunk != nil; chunk = chunk.head.next { 716 | if chunk.Type == Head { 717 | if chunkHead != nil { 718 | return errors.New("multiple chunk heads found") 719 | } 720 | chunkHead = chunk 721 | } 722 | } 723 | if chunkHead != chunks.first { 724 | return errors.New("chunk head does not have the newest timestamp") 725 | } 726 | return nil 727 | } 728 | 729 | func scanJournals(logger *logging.Logger, pathPrefix string, pathSuffix string) (map[string]*FileJournal, error) { 730 | journals := make(map[string]*FileJournal) 731 | dirname, basename := filepath.Split(pathPrefix) 732 | if dirname == "" { 733 | dirname = "." 734 | } 735 | d, err := os.OpenFile(dirname, os.O_RDONLY, 0) 736 | if err != nil { 737 | return nil, err 738 | } 739 | defer d.Close() 740 | finfo, err := d.Stat() 741 | if err != nil { 742 | return nil, err 743 | } 744 | if !finfo.IsDir() { 745 | return nil, errors.New(fmt.Sprintf("%s is not a directory", dirname)) 746 | } 747 | for { 748 | files_, err := d.Readdir(100) 749 | if err == io.EOF { 750 | break 751 | } else if err != nil { 752 | return nil, err 753 | } 754 | for _, finfo := range files_ { 755 | file := finfo.Name() 756 | if !strings.HasSuffix(file, pathSuffix) { 757 | continue 758 | } 759 | variablePortion := file[len(basename) : len(file)-len(pathSuffix)] 760 | info, err := DecodeJournalPath(variablePortion) 761 | if err != nil { 762 | logger.Warningf("Unexpected file under the designated directory space (%s) - %s", dirname, file) 763 | continue 764 | } 765 | journalProto, ok := journals[info.Key] 766 | if !ok { 767 | journalProto = &FileJournal{ 768 | key: info.Key, 769 | chunks: FileJournalChunkDequeue{nil, nil, 0, sync.Mutex{}}, 770 | writer: nil, 771 | } 772 | journals[info.Key] = journalProto 773 | } 774 | chunk := &FileJournalChunk{ 775 | head: FileJournalChunkDequeueHead{nil, journalProto.chunks.last}, 776 | container: &journalProto.chunks, 777 | Type: info.Type, 778 | Path: pathPrefix + info.VariablePortion + pathSuffix, 779 | TSuffix: info.TSuffix, 780 | Timestamp: info.Timestamp, 781 | UniqueId: info.UniqueId, 782 | Size: finfo.Size(), 783 | refcount: 1, 784 | } 785 | if journalProto.chunks.last == nil { 786 | journalProto.chunks.first = chunk 787 | } else { 788 | journalProto.chunks.last.head.next = chunk 789 | } 790 | journalProto.chunks.last = chunk 791 | journalProto.chunks.count += 1 792 | } 793 | } 794 | for _, journalProto := range journals { 795 | sortChunksByTimestamp(&journalProto.chunks) 796 | err := validateChunks(&journalProto.chunks) 797 | if err != nil { 798 | return nil, err 799 | } 800 | } 801 | return journals, nil 802 | } 803 | 804 | func (factory *FileJournalGroupFactory) GetJournalGroup(path string, worker Worker) (*FileJournalGroup, error) { 805 | registered, ok := factory.paths[path] 806 | if ok { 807 | if registered.worker == worker { 808 | return registered, nil 809 | } else { 810 | return nil, errors.New(fmt.Sprintf( 811 | "Other worker '%s' already use same buffer_path: %s", 812 | registered.worker.String(), 813 | path, 814 | )) 815 | } 816 | } 817 | 818 | var pathPrefix string 819 | var pathSuffix string 820 | 821 | pos := strings.Index(path, "*") 822 | if pos >= 0 { 823 | pathPrefix = path[0:pos] 824 | pathSuffix = path[pos+1:] 825 | } else { 826 | pathPrefix = path + "." 827 | pathSuffix = factory.defaultPathSuffix 828 | } 829 | 830 | journals, err := scanJournals(factory.logger, pathPrefix, pathSuffix) 831 | if err != nil { 832 | return nil, err 833 | } 834 | 835 | journalGroup := &FileJournalGroup{ 836 | factory: factory, 837 | worker: worker, 838 | timeGetter: factory.timeGetter, 839 | logger: factory.logger, 840 | rand: rand.New(factory.randSource), 841 | fileMode: factory.defaultFileMode, 842 | maxSize: factory.maxSize, 843 | pathPrefix: pathPrefix, 844 | pathSuffix: pathSuffix, 845 | journals: journals, 846 | mtx: sync.Mutex{}, 847 | } 848 | for _, journal := range journals { 849 | journal.group = journalGroup 850 | journal.newChunkListeners = make(map[JournalChunkListener]JournalChunkListener) 851 | journal.flushListeners = make(map[JournalChunkListener]JournalChunkListener) 852 | chunk := journal.chunks.first 853 | file, err := os.OpenFile(chunk.Path, os.O_WRONLY|os.O_APPEND, journal.group.fileMode) 854 | if err != nil { 855 | journalGroup.Dispose() 856 | return nil, err 857 | } 858 | position, err := file.Seek(0, os.SEEK_END) 859 | if err != nil { 860 | file.Close() 861 | journalGroup.Dispose() 862 | return nil, err 863 | } 864 | chunk.refcount += 1 // for writer 865 | chunk.Size = position 866 | journal.writer = file 867 | } 868 | factory.logger.Infof("Path %s is designated to Worker %s", path, worker.String()) 869 | factory.paths[path] = journalGroup 870 | return journalGroup, nil 871 | } 872 | 873 | func NewFileJournalGroupFactory( 874 | logger *logging.Logger, 875 | randSource rand.Source, 876 | timeGetter func() time.Time, 877 | defaultPathSuffix string, 878 | defaultFileMode os.FileMode, 879 | maxSize int64, 880 | ) *FileJournalGroupFactory { 881 | return &FileJournalGroupFactory{ 882 | logger: logger, 883 | paths: make(map[string]*FileJournalGroup), 884 | randSource: randSource, 885 | timeGetter: timeGetter, 886 | defaultPathSuffix: defaultPathSuffix, 887 | defaultFileMode: defaultFileMode, 888 | maxSize: maxSize, 889 | } 890 | } 891 | -------------------------------------------------------------------------------- /file_journal_test.go: -------------------------------------------------------------------------------- 1 | // 2 | // Copyright (C) 2014 Moriyoshi Koizumi 3 | // Copyright (C) 2014 Treasure Data, Inc. 4 | // 5 | // Permission is hereby granted, free of charge, to any person obtaining a copy 6 | // of this software and associated documentation files (the "Software"), to deal 7 | // in the Software without restriction, including without limitation the rights 8 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | // copies of the Software, and to permit persons to whom the Software is 10 | // furnished to do so, subject to the following conditions: 11 | // 12 | // The above copyright notice and this permission notice shall be included in 13 | // all copies or substantial portions of the Software. 14 | // 15 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 21 | // THE SOFTWARE. 22 | // 23 | 24 | package fluentd_forwarder 25 | 26 | import ( 27 | "bufio" 28 | "fmt" 29 | logging "github.com/op/go-logging" 30 | "io" 31 | "io/ioutil" 32 | "math/rand" 33 | "os" 34 | "path/filepath" 35 | "runtime" 36 | "sync" 37 | "sync/atomic" 38 | "testing" 39 | "time" 40 | ) 41 | 42 | type DummyWorker struct{ v int } 43 | 44 | func (*DummyWorker) String() string { return "" } 45 | func (*DummyWorker) Start() {} 46 | func (*DummyWorker) Stop() {} 47 | func (*DummyWorker) WaitForShutdown() {} 48 | 49 | type DummyChunkListener struct { 50 | t *testing.T 51 | chunks []FileJournalChunk 52 | } 53 | 54 | func (*DummyChunkListener) NewChunkCreated(chunk JournalChunk) error { 55 | return nil 56 | } 57 | 58 | func (listener *DummyChunkListener) ChunkFlushed(chunk JournalChunk) error { 59 | defer chunk.Dispose() 60 | impl := chunk.(*FileJournalChunkWrapper) 61 | listener.chunks = append(listener.chunks, *impl.chunk) 62 | listener.t.Logf("flush %d", len(listener.chunks)) 63 | return nil 64 | } 65 | 66 | func Test_GetJournalGroup(t *testing.T) { 67 | logging.InitForTesting(logging.NOTICE) 68 | logger := logging.MustGetLogger("journal") 69 | tempDir, err := ioutil.TempDir("", "journal") 70 | if err != nil { 71 | t.FailNow() 72 | } 73 | defer os.RemoveAll(tempDir) 74 | factory := NewFileJournalGroupFactory( 75 | logger, 76 | rand.NewSource(0), 77 | func() time.Time { return time.Date(2014, 1, 1, 0, 0, 0, 0, time.UTC) }, 78 | ".log", 79 | os.FileMode(0644), 80 | 0, 81 | ) 82 | dummyWorker := &DummyWorker{} 83 | tempFile := filepath.Join(tempDir, "test") 84 | t.Log(tempFile) 85 | journalGroup1, err := factory.GetJournalGroup(tempFile, dummyWorker) 86 | if err != nil { 87 | t.FailNow() 88 | } 89 | journalGroup2, err := factory.GetJournalGroup(tempFile, dummyWorker) 90 | if err != nil { 91 | t.Fail() 92 | } 93 | if journalGroup1 != journalGroup2 { 94 | t.Fail() 95 | } 96 | anotherDummyWorker := &DummyWorker{} 97 | if dummyWorker == anotherDummyWorker { 98 | t.Log("WTF?") 99 | t.Fail() 100 | } 101 | _, err = factory.GetJournalGroup(tempFile, anotherDummyWorker) 102 | if err == nil { 103 | t.Fail() 104 | } 105 | } 106 | 107 | func Test_Journal_GetJournal(t *testing.T) { 108 | logging.InitForTesting(logging.NOTICE) 109 | logger := logging.MustGetLogger("journal") 110 | tempDir, err := ioutil.TempDir("", "journal") 111 | if err != nil { 112 | t.FailNow() 113 | } 114 | defer os.RemoveAll(tempDir) 115 | factory := NewFileJournalGroupFactory( 116 | logger, 117 | rand.NewSource(0), 118 | func() time.Time { return time.Date(2014, 1, 1, 0, 0, 0, 0, time.UTC) }, 119 | ".log", 120 | os.FileMode(0644), 121 | 0, 122 | ) 123 | dummyWorker := &DummyWorker{} 124 | tempFile := filepath.Join(tempDir, "test") 125 | t.Log(tempFile) 126 | journalGroup, err := factory.GetJournalGroup(tempFile, dummyWorker) 127 | if err != nil { 128 | t.FailNow() 129 | } 130 | journal1 := journalGroup.GetJournal("key") 131 | if journal1 == nil { 132 | t.FailNow() 133 | } 134 | journal2 := journalGroup.GetJournal("key") 135 | if journal2 == nil { 136 | t.Fail() 137 | } 138 | if journal1 != journal2 { 139 | t.Fail() 140 | } 141 | } 142 | 143 | func Test_Journal_EmitVeryFirst(t *testing.T) { 144 | logging.InitForTesting(logging.NOTICE) 145 | logger := logging.MustGetLogger("journal") 146 | tempDir, err := ioutil.TempDir("", "journal") 147 | if err != nil { 148 | t.FailNow() 149 | } 150 | defer os.RemoveAll(tempDir) 151 | factory := NewFileJournalGroupFactory( 152 | logger, 153 | rand.NewSource(0), 154 | func() time.Time { return time.Date(2014, 1, 1, 0, 0, 0, 0, time.UTC) }, 155 | ".log", 156 | os.FileMode(0644), 157 | 10, 158 | ) 159 | dummyWorker := &DummyWorker{} 160 | tempFile := filepath.Join(tempDir, "test") 161 | t.Log(tempFile) 162 | journalGroup, err := factory.GetJournalGroup(tempFile, dummyWorker) 163 | if err != nil { 164 | t.FailNow() 165 | } 166 | journal := journalGroup.GetFileJournal("key") 167 | defer journal.Dispose() 168 | err = journal.Write([]byte("test")) 169 | if err != nil { 170 | t.FailNow() 171 | } 172 | if journal.chunks.count != 1 { 173 | t.Fail() 174 | } 175 | if journal.chunks.first.Size != 4 { 176 | t.Fail() 177 | } 178 | } 179 | 180 | func Test_Journal_EmitTwice(t *testing.T) { 181 | logging.InitForTesting(logging.NOTICE) 182 | logger := logging.MustGetLogger("journal") 183 | tempDir, err := ioutil.TempDir("", "journal") 184 | if err != nil { 185 | t.FailNow() 186 | } 187 | defer os.RemoveAll(tempDir) 188 | factory := NewFileJournalGroupFactory( 189 | logger, 190 | rand.NewSource(0), 191 | func() time.Time { return time.Date(2014, 1, 1, 0, 0, 0, 0, time.UTC) }, 192 | ".log", 193 | os.FileMode(0644), 194 | 10, 195 | ) 196 | dummyWorker := &DummyWorker{} 197 | tempFile := filepath.Join(tempDir, "test") 198 | t.Log(tempFile) 199 | journalGroup, err := factory.GetJournalGroup(tempFile, dummyWorker) 200 | if err != nil { 201 | t.FailNow() 202 | } 203 | journal := journalGroup.GetFileJournal("key") 204 | defer journal.Dispose() 205 | err = journal.Write([]byte("test1")) 206 | if err != nil { 207 | t.FailNow() 208 | } 209 | err = journal.Write([]byte("test2")) 210 | if err != nil { 211 | t.FailNow() 212 | } 213 | if journal.chunks.count != 1 { 214 | t.Fail() 215 | } 216 | if journal.chunks.first.Size != 10 { 217 | t.Fail() 218 | } 219 | } 220 | 221 | func Test_Journal_EmitRotating(t *testing.T) { 222 | logging.InitForTesting(logging.NOTICE) 223 | logger := logging.MustGetLogger("journal") 224 | tempDir, err := ioutil.TempDir("", "journal") 225 | if err != nil { 226 | t.FailNow() 227 | } 228 | defer os.RemoveAll(tempDir) 229 | factory := NewFileJournalGroupFactory( 230 | logger, 231 | rand.NewSource(0), 232 | func() time.Time { return time.Date(2014, 1, 1, 0, 0, 0, 0, time.UTC) }, 233 | ".log", 234 | os.FileMode(0644), 235 | 8, 236 | ) 237 | dummyWorker := &DummyWorker{} 238 | tempFile := filepath.Join(tempDir, "test") 239 | t.Log(tempFile) 240 | journalGroup, err := factory.GetJournalGroup(tempFile, dummyWorker) 241 | if err != nil { 242 | t.FailNow() 243 | } 244 | journal := journalGroup.GetFileJournal("key") 245 | defer journal.Dispose() 246 | err = journal.Write([]byte("test1")) 247 | if err != nil { 248 | t.FailNow() 249 | } 250 | err = journal.Write([]byte("test2")) 251 | if err != nil { 252 | t.FailNow() 253 | } 254 | err = journal.Write([]byte("test3")) 255 | if err != nil { 256 | t.FailNow() 257 | } 258 | err = journal.Write([]byte("test4")) 259 | if err != nil { 260 | t.FailNow() 261 | } 262 | err = journal.Write([]byte("test5")) 263 | if err != nil { 264 | t.FailNow() 265 | } 266 | t.Logf("journal.chunks.count=%d", journal.chunks.count) 267 | t.Logf("journal.chunks.first.Size=%d", journal.chunks.first.Size) 268 | if journal.chunks.count != 5 { 269 | t.Fail() 270 | } 271 | if journal.chunks.first.Size != 5 { 272 | t.Fail() 273 | } 274 | } 275 | 276 | func shuffle(x []string) { 277 | rng := rand.New(rand.NewSource(0)) 278 | for i := 0; i < len(x); i += 1 { 279 | j := rng.Intn(i + 1) 280 | x[i], x[j] = x[j], x[i] 281 | } 282 | } 283 | 284 | func Test_Journal_Scanning_Ok(t *testing.T) { 285 | logging.InitForTesting(logging.NOTICE) 286 | logger := logging.MustGetLogger("journal") 287 | tm := time.Date(2014, 1, 1, 0, 0, 0, 0, time.UTC) 288 | 289 | for i := 1; i < 100; i++ { 290 | tempDir, err := ioutil.TempDir("", "journal") 291 | if err != nil { 292 | t.FailNow() 293 | } 294 | defer os.RemoveAll(tempDir) 295 | prefix := filepath.Join(tempDir, "test") 296 | suffix := ".log" 297 | makePaths := func(n int, key string) []string { 298 | paths := make([]string, n) 299 | for i := 0; i < len(paths); i += 1 { 300 | type_ := JournalFileType('q') 301 | if i == 0 { 302 | type_ = JournalFileType('b') 303 | } 304 | path := prefix + "." + BuildJournalPath(key, type_, tm.Add(time.Duration(-i*1e9)), 0).VariablePortion + suffix 305 | paths[i] = path 306 | } 307 | return paths 308 | } 309 | 310 | paths := makePaths(i, "key") 311 | shuffledPaths := make([]string, len(paths)) 312 | copy(shuffledPaths, paths) 313 | shuffle(shuffledPaths) 314 | for j, path := range shuffledPaths { 315 | file, err := os.Create(path) 316 | if err != nil { 317 | t.FailNow() 318 | } 319 | _, err = file.Write([]byte(fmt.Sprintf("%08d", j))) 320 | if err != nil { 321 | t.FailNow() 322 | } 323 | file.Close() 324 | } 325 | factory := NewFileJournalGroupFactory( 326 | logger, 327 | rand.NewSource(0), 328 | func() time.Time { return tm }, 329 | suffix, 330 | os.FileMode(0644), 331 | 8, 332 | ) 333 | dummyWorker := &DummyWorker{} 334 | journalGroup, err := factory.GetJournalGroup(prefix, dummyWorker) 335 | if err != nil { 336 | t.FailNow() 337 | } 338 | journal := journalGroup.GetFileJournal("key") 339 | defer journal.Dispose() 340 | t.Logf("journal.chunks.count=%d", journal.chunks.count) 341 | t.Logf("journal.chunks.first.Size=%d", journal.chunks.first.Size) 342 | if journal.chunks.count != i { 343 | t.Fail() 344 | } 345 | j := 0 346 | for chunk := journal.chunks.first; chunk != nil; chunk = chunk.head.next { 347 | if chunk.Path != paths[j] { 348 | t.Fail() 349 | } 350 | j += 1 351 | } 352 | journal.Flush(nil) 353 | os.RemoveAll(tempDir) 354 | t.Logf("journal.chunks.count=%d", journal.chunks.count) 355 | if journal.chunks.count != 1 { 356 | t.Fail() 357 | } 358 | } 359 | } 360 | 361 | func Test_Journal_Scanning_MultipleHead(t *testing.T) { 362 | logging.InitForTesting(logging.NOTICE) 363 | logger := logging.MustGetLogger("journal") 364 | tempDir, err := ioutil.TempDir("", "journal") 365 | if err != nil { 366 | t.FailNow() 367 | } 368 | defer os.RemoveAll(tempDir) 369 | tm := time.Date(2014, 1, 1, 0, 0, 0, 0, time.UTC) 370 | prefix := filepath.Join(tempDir, "test") 371 | suffix := ".log" 372 | createFile := func(key string, type_ JournalFileType, o int) (string, error) { 373 | path := prefix + "." + BuildJournalPath(key, type_, tm.Add(time.Duration(-o*1e9)), 0).VariablePortion + suffix 374 | file, err := os.Create(path) 375 | if err != nil { 376 | return "", err 377 | } 378 | _, err = file.Write([]byte(fmt.Sprintf("%08d", o))) 379 | if err != nil { 380 | return "", err 381 | } 382 | file.Close() 383 | t.Log(path) 384 | return path, nil 385 | } 386 | 387 | paths := make([]string, 4) 388 | { 389 | path, err := createFile("key", JournalFileType('b'), 0) 390 | if err != nil { 391 | t.FailNow() 392 | } 393 | paths[0] = path 394 | } 395 | { 396 | path, err := createFile("key", JournalFileType('b'), 1) 397 | if err != nil { 398 | t.FailNow() 399 | } 400 | paths[1] = path 401 | } 402 | { 403 | path, err := createFile("key", JournalFileType('q'), 2) 404 | if err != nil { 405 | t.FailNow() 406 | } 407 | paths[2] = path 408 | } 409 | { 410 | path, err := createFile("key", JournalFileType('q'), 3) 411 | if err != nil { 412 | t.FailNow() 413 | } 414 | paths[3] = path 415 | } 416 | 417 | factory := NewFileJournalGroupFactory( 418 | logger, 419 | rand.NewSource(0), 420 | func() time.Time { return tm }, 421 | suffix, 422 | os.FileMode(0644), 423 | 8, 424 | ) 425 | dummyWorker := &DummyWorker{} 426 | _, err = factory.GetJournalGroup(prefix, dummyWorker) 427 | if err == nil { 428 | t.FailNow() 429 | } 430 | t.Log(err.Error()) 431 | if err.Error() != "multiple chunk heads found" { 432 | t.Fail() 433 | } 434 | } 435 | 436 | func Test_Journal_FlushListener(t *testing.T) { 437 | logging.InitForTesting(logging.NOTICE) 438 | logger := logging.MustGetLogger("journal") 439 | tempDir, err := ioutil.TempDir("", "journal") 440 | if err != nil { 441 | t.FailNow() 442 | } 443 | defer os.RemoveAll(tempDir) 444 | factory := NewFileJournalGroupFactory( 445 | logger, 446 | rand.NewSource(0), 447 | func() time.Time { return time.Date(2014, 1, 1, 0, 0, 0, 0, time.UTC) }, 448 | ".log", 449 | os.FileMode(0644), 450 | 8, 451 | ) 452 | dummyWorker := &DummyWorker{} 453 | tempFile := filepath.Join(tempDir, "test") 454 | t.Log(tempFile) 455 | journalGroup, err := factory.GetJournalGroup(tempFile, dummyWorker) 456 | if err != nil { 457 | t.FailNow() 458 | } 459 | journal := journalGroup.GetFileJournal("key") 460 | defer journal.Dispose() 461 | listener := &DummyChunkListener{ 462 | t: t, 463 | chunks: make([]FileJournalChunk, 0, 5), 464 | } 465 | journal.AddFlushListener(listener) 466 | journal.AddFlushListener(listener) 467 | err = journal.Write([]byte("test1")) 468 | if err != nil { 469 | t.FailNow() 470 | } 471 | err = journal.Write([]byte("test2")) 472 | if err != nil { 473 | t.FailNow() 474 | } 475 | err = journal.Write([]byte("test3")) 476 | if err != nil { 477 | t.FailNow() 478 | } 479 | err = journal.Write([]byte("test4")) 480 | if err != nil { 481 | t.FailNow() 482 | } 483 | err = journal.Write([]byte("test5")) 484 | if err != nil { 485 | t.FailNow() 486 | } 487 | t.Logf("journal.chunks.count=%d", journal.chunks.count) 488 | t.Logf("journal.chunks.first.Size=%d", journal.chunks.first.Size) 489 | t.Logf("len(listener.chunks)=%d", len(listener.chunks)) 490 | if journal.chunks.count != 5 { 491 | t.Fail() 492 | } 493 | if journal.chunks.first.Size != 5 { 494 | t.Fail() 495 | } 496 | if len(listener.chunks) != 4 { 497 | t.Fail() 498 | } 499 | readAll := func(chunk FileJournalChunk) string { 500 | reader, err := chunk.getReader() 501 | if err != nil { 502 | t.FailNow() 503 | } 504 | defer reader.Close() 505 | bytes, err := ioutil.ReadAll(reader) 506 | if err != nil { 507 | t.FailNow() 508 | } 509 | return string(bytes) 510 | } 511 | if readAll(listener.chunks[0]) != "test1" { 512 | t.Fail() 513 | } 514 | if readAll(listener.chunks[1]) != "test2" { 515 | t.Fail() 516 | } 517 | if readAll(listener.chunks[2]) != "test3" { 518 | t.Fail() 519 | } 520 | if readAll(listener.chunks[3]) != "test4" { 521 | t.Fail() 522 | } 523 | journal.Flush(nil) 524 | t.Logf("journal.chunks.count=%d", journal.chunks.count) 525 | t.Logf("journal.chunks.first.Size=%d", journal.chunks.first.Size) 526 | if journal.chunks.count != 1 { 527 | t.Fail() 528 | } 529 | if journal.chunks.first.Type != JournalFileType('b') { 530 | t.Fail() 531 | } 532 | } 533 | 534 | func countLines(r io.Reader) (int, error) { 535 | count := 0 536 | bi := bufio.NewReader(r) 537 | for { 538 | _, _, err := bi.ReadLine() 539 | if err == io.EOF { 540 | break 541 | } else if err != nil { 542 | return 0, err 543 | } 544 | count += 1 545 | } 546 | return count, nil 547 | } 548 | 549 | func Test_Journal_Concurrency(t *testing.T) { 550 | logging.InitForTesting(logging.NOTICE) 551 | logger := logging.MustGetLogger("journal") 552 | tempDir, err := ioutil.TempDir("", "journal") 553 | if err != nil { 554 | t.FailNow() 555 | } 556 | defer os.RemoveAll(tempDir) 557 | factory := NewFileJournalGroupFactory( 558 | logger, 559 | rand.NewSource(0), 560 | func() time.Time { return time.Date(2014, 1, 1, 0, 0, 0, 0, time.UTC) }, 561 | ".log", 562 | os.FileMode(0644), 563 | 16, 564 | ) 565 | dummyWorker := &DummyWorker{} 566 | tempFile := filepath.Join(tempDir, "test") 567 | t.Log(tempFile) 568 | journalGroup, err := factory.GetJournalGroup(tempFile, dummyWorker) 569 | if err != nil { 570 | t.FailNow() 571 | } 572 | journal := journalGroup.GetFileJournal("key") 573 | defer journal.Dispose() 574 | cond := sync.Cond{L: &sync.Mutex{}} 575 | count := int64(0) 576 | outerWg := sync.WaitGroup{} 577 | doFlush := func() { 578 | journal.Flush(func(chunk JournalChunk) interface{} { 579 | defer chunk.Dispose() 580 | reader, err := chunk.Reader() 581 | if err != nil { 582 | t.Log(err.Error()) 583 | t.FailNow() 584 | } 585 | defer reader.Close() 586 | c, err := countLines(reader) 587 | if err != nil { 588 | t.Log(err.Error()) 589 | t.FailNow() 590 | } 591 | atomic.AddInt64(&count, int64(c)) 592 | return nil 593 | }) 594 | } 595 | for j := 0; j < 10; j += 1 { 596 | outerWg.Add(1) 597 | go func(j int) { 598 | defer outerWg.Done() 599 | wg := sync.WaitGroup{} 600 | starting := sync.WaitGroup{} 601 | for i := 0; i < 10; i += 1 { 602 | wg.Add(1) 603 | starting.Add(1) 604 | go func(i int) { 605 | defer wg.Done() 606 | starting.Done() 607 | cond.L.Lock() 608 | cond.Wait() 609 | cond.L.Unlock() 610 | for k := 0; k < 3; k += 1 { 611 | data := fmt.Sprintf("test%d\n", i) 612 | err = journal.Write([]byte(data)) 613 | if err != nil { 614 | t.Log(err.Error()) 615 | t.FailNow() 616 | } 617 | } 618 | }(i) 619 | } 620 | starting.Wait() 621 | cond.Broadcast() 622 | runtime.Gosched() 623 | doFlush() 624 | wg.Wait() 625 | }(j) 626 | } 627 | outerWg.Wait() 628 | doFlush() 629 | if count != 300 { 630 | t.Logf("%d", count) 631 | t.Fail() 632 | } 633 | } 634 | -------------------------------------------------------------------------------- /forwarder.go: -------------------------------------------------------------------------------- 1 | // 2 | // Fluentd Forwarder 3 | // 4 | // Copyright (C) 2014 Treasure Data, Inc. 5 | // 6 | // Licensed under the Apache License, Version 2.0 (the "License"); 7 | // you may not use this file except in compliance with the License. 8 | // You may obtain a copy of the License at 9 | // 10 | // http://www.apache.org/licenses/LICENSE-2.0 11 | // 12 | // Unless required by applicable law or agreed to in writing, software 13 | // distributed under the License is distributed on an "AS IS" BASIS, 14 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | // See the License for the specific language governing permissions and 16 | // limitations under the License. 17 | // 18 | 19 | package fluentd_forwarder 20 | 21 | import ( 22 | "fmt" 23 | "io" 24 | ) 25 | 26 | type FluentRecord struct { 27 | Tag string 28 | Timestamp uint64 29 | Data map[string]interface{} 30 | } 31 | 32 | type TinyFluentRecord struct { 33 | Timestamp uint64 34 | Data map[string]interface{} 35 | } 36 | 37 | type FluentRecordSet struct { 38 | Tag string 39 | Records []TinyFluentRecord 40 | } 41 | 42 | type Port interface { 43 | Emit(recordSets []FluentRecordSet) error 44 | } 45 | 46 | type Worker interface { 47 | String() string 48 | Start() 49 | Stop() 50 | WaitForShutdown() 51 | } 52 | 53 | type Disposable interface { 54 | Dispose() error 55 | } 56 | 57 | type JournalChunk interface { 58 | Disposable 59 | Id() string 60 | String() string 61 | Size() (int64, error) 62 | Reader() (io.ReadCloser, error) 63 | NextChunk() JournalChunk 64 | MD5Sum() ([]byte, error) 65 | Dup() JournalChunk 66 | } 67 | 68 | type JournalChunkListener interface { 69 | NewChunkCreated(JournalChunk) error 70 | ChunkFlushed(JournalChunk) error 71 | } 72 | 73 | type Journal interface { 74 | Disposable 75 | Key() string 76 | Write(data []byte) error 77 | TailChunk() JournalChunk 78 | AddNewChunkListener(JournalChunkListener) 79 | AddFlushListener(JournalChunkListener) 80 | Flush(func(JournalChunk) interface{}) error 81 | } 82 | 83 | type JournalGroup interface { 84 | Disposable 85 | GetJournal(key string) Journal 86 | GetJournalKeys() []string 87 | } 88 | 89 | type JournalGroupFactory interface { 90 | GetJournalGroup() JournalGroup 91 | } 92 | 93 | type Panicked struct { 94 | Opaque interface{} 95 | } 96 | 97 | func (e *Panicked) Error() string { 98 | s, ok := e.Opaque.(string) 99 | if ok { 100 | return s 101 | } 102 | return fmt.Sprintf("%v", e.Opaque) 103 | } 104 | -------------------------------------------------------------------------------- /input.go: -------------------------------------------------------------------------------- 1 | // 2 | // Fluentd Forwarder 3 | // 4 | // Copyright (C) 2014 Treasure Data, Inc. 5 | // 6 | // Licensed under the Apache License, Version 2.0 (the "License"); 7 | // you may not use this file except in compliance with the License. 8 | // You may obtain a copy of the License at 9 | // 10 | // http://www.apache.org/licenses/LICENSE-2.0 11 | // 12 | // Unless required by applicable law or agreed to in writing, software 13 | // distributed under the License is distributed on an "AS IS" BASIS, 14 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | // See the License for the specific language governing permissions and 16 | // limitations under the License. 17 | // 18 | 19 | package fluentd_forwarder 20 | 21 | import ( 22 | "bufio" 23 | "bytes" 24 | "errors" 25 | "fmt" 26 | logging "github.com/op/go-logging" 27 | "github.com/ugorji/go/codec" 28 | "io" 29 | "net" 30 | "reflect" 31 | "sync" 32 | "sync/atomic" 33 | ) 34 | 35 | type forwardClient struct { 36 | input *ForwardInput 37 | logger *logging.Logger 38 | conn *net.TCPConn 39 | codec *codec.MsgpackHandle 40 | dec *codec.Decoder 41 | } 42 | 43 | type ForwardInput struct { 44 | entries int64 // This variable must be on 64-bit alignment. Otherwise atomic.AddInt64 will cause a crash on ARM and x86-32 45 | port Port 46 | logger *logging.Logger 47 | bind string 48 | listener *net.TCPListener 49 | codec *codec.MsgpackHandle 50 | clientsMtx sync.Mutex 51 | clients map[*net.TCPConn]*forwardClient 52 | wg sync.WaitGroup 53 | acceptChan chan *net.TCPConn 54 | shutdownChan chan struct{} 55 | isShuttingDown uintptr 56 | } 57 | 58 | type EntryCountTopic struct{} 59 | 60 | type ConnectionCountTopic struct{} 61 | 62 | type ForwardInputFactory struct{} 63 | 64 | func coerceInPlace(data map[string]interface{}) { 65 | for k, v := range data { 66 | switch v_ := v.(type) { 67 | case []byte: 68 | data[k] = string(v_) // XXX: byte => rune 69 | case map[string]interface{}: 70 | coerceInPlace(v_) 71 | } 72 | } 73 | } 74 | 75 | func (c *forwardClient) decodeRecordSet(tag []byte, entries []interface{}) (FluentRecordSet, error) { 76 | records := make([]TinyFluentRecord, len(entries)) 77 | for i, _entry := range entries { 78 | entry, ok := _entry.([]interface{}) 79 | if !ok { 80 | return FluentRecordSet{}, errors.New("Failed to decode recordSet") 81 | } 82 | timestamp, ok := entry[0].(uint64) 83 | if !ok { 84 | return FluentRecordSet{}, errors.New("Failed to decode timestamp field") 85 | } 86 | data, ok := entry[1].(map[string]interface{}) 87 | if !ok { 88 | return FluentRecordSet{}, errors.New("Failed to decode data field") 89 | } 90 | coerceInPlace(data) 91 | records[i] = TinyFluentRecord{ 92 | Timestamp: timestamp, 93 | Data: data, 94 | } 95 | } 96 | return FluentRecordSet{ 97 | Tag: string(tag), // XXX: byte => rune 98 | Records: records, 99 | }, nil 100 | } 101 | 102 | func (c *forwardClient) decodeEntries() ([]FluentRecordSet, error) { 103 | v := []interface{}{nil, nil, nil} 104 | err := c.dec.Decode(&v) 105 | if err != nil { 106 | return nil, err 107 | } 108 | tag, ok := v[0].([]byte) 109 | if !ok { 110 | return nil, errors.New("Failed to decode tag field") 111 | } 112 | 113 | var retval []FluentRecordSet 114 | switch timestamp_or_entries := v[1].(type) { 115 | case uint64: 116 | timestamp := timestamp_or_entries 117 | data, ok := v[2].(map[string]interface{}) 118 | if !ok { 119 | return nil, errors.New("Failed to decode data field") 120 | } 121 | coerceInPlace(data) 122 | retval = []FluentRecordSet{ 123 | { 124 | Tag: string(tag), // XXX: byte => rune 125 | Records: []TinyFluentRecord{ 126 | { 127 | Timestamp: timestamp, 128 | Data: data, 129 | }, 130 | }, 131 | }, 132 | } 133 | case float64: 134 | timestamp := uint64(timestamp_or_entries) 135 | data, ok := v[2].(map[string]interface{}) 136 | if !ok { 137 | return nil, errors.New("Failed to decode data field") 138 | } 139 | retval = []FluentRecordSet{ 140 | { 141 | Tag: string(tag), // XXX: byte => rune 142 | Records: []TinyFluentRecord{ 143 | { 144 | Timestamp: timestamp, 145 | Data: data, 146 | }, 147 | }, 148 | }, 149 | } 150 | case []interface{}: 151 | if !ok { 152 | return nil, errors.New("Unexpected payload format") 153 | } 154 | recordSet, err := c.decodeRecordSet(tag, timestamp_or_entries) 155 | if err != nil { 156 | return nil, err 157 | } 158 | retval = []FluentRecordSet{recordSet} 159 | case []byte: 160 | entries := make([]interface{}, 0) 161 | reader := bytes.NewReader(timestamp_or_entries) 162 | dec := codec.NewDecoder(reader, c.codec) 163 | for reader.Len() > 0 { // codec.Decoder doesn't return EOF. 164 | entry := []interface{}{} 165 | if err != dec.Decode(&entry) { 166 | if err == io.EOF { // in case codec.Decoder changes its behavior 167 | break 168 | } 169 | return nil, err 170 | } 171 | entries = append(entries, entry) 172 | } 173 | recordSet, err := c.decodeRecordSet(tag, entries) 174 | if err != nil { 175 | return nil, err 176 | } 177 | retval = []FluentRecordSet{recordSet} 178 | default: 179 | return nil, errors.New(fmt.Sprintf("Unknown type: %t", timestamp_or_entries)) 180 | } 181 | atomic.AddInt64(&c.input.entries, int64(len(retval))) 182 | return retval, nil 183 | } 184 | 185 | func (c *forwardClient) startHandling() { 186 | c.input.wg.Add(1) 187 | go func() { 188 | defer func() { 189 | err := c.conn.Close() 190 | if err != nil { 191 | c.logger.Debugf("Close: %s", err.Error()) 192 | } 193 | c.input.markDischarged(c) 194 | c.input.wg.Done() 195 | }() 196 | c.input.logger.Infof("Started handling connection from %s", c.conn.RemoteAddr().String()) 197 | for { 198 | recordSets, err := c.decodeEntries() 199 | if err != nil { 200 | err_, ok := err.(net.Error) 201 | if ok { 202 | if err_.Temporary() { 203 | c.logger.Infof("Temporary failure: %s", err_.Error()) 204 | continue 205 | } 206 | } 207 | if err == io.EOF { 208 | c.logger.Infof("Client %s closed the connection", c.conn.RemoteAddr().String()) 209 | } else { 210 | c.logger.Error(err.Error()) 211 | } 212 | break 213 | } 214 | 215 | if len(recordSets) > 0 { 216 | err_ := c.input.port.Emit(recordSets) 217 | if err_ != nil { 218 | c.logger.Error(err_.Error()) 219 | break 220 | } 221 | } 222 | } 223 | c.input.logger.Infof("Ended handling connection from %s", c.conn.RemoteAddr().String()) 224 | }() 225 | } 226 | 227 | func (c *forwardClient) shutdown() { 228 | err := c.conn.Close() 229 | if err != nil { 230 | c.input.logger.Infof("Error during closing connection: %s", err.Error()) 231 | } 232 | } 233 | 234 | func newForwardClient(input *ForwardInput, logger *logging.Logger, conn *net.TCPConn, _codec *codec.MsgpackHandle) *forwardClient { 235 | c := &forwardClient{ 236 | input: input, 237 | logger: logger, 238 | conn: conn, 239 | codec: _codec, 240 | dec: codec.NewDecoder(bufio.NewReader(conn), _codec), 241 | } 242 | input.markCharged(c) 243 | return c 244 | } 245 | 246 | func (input *ForwardInput) spawnAcceptor() { 247 | input.logger.Notice("Spawning acceptor") 248 | input.wg.Add(1) 249 | go func() { 250 | defer func() { 251 | close(input.acceptChan) 252 | input.wg.Done() 253 | }() 254 | input.logger.Notice("Acceptor started") 255 | for { 256 | conn, err := input.listener.AcceptTCP() 257 | if err != nil { 258 | input.logger.Notice(err.Error()) 259 | break 260 | } 261 | if conn != nil { 262 | input.logger.Noticef("Connected from %s", conn.RemoteAddr().String()) 263 | input.acceptChan <- conn 264 | } else { 265 | input.logger.Notice("AcceptTCP returned nil; something went wrong") 266 | break 267 | } 268 | } 269 | input.logger.Notice("Acceptor ended") 270 | }() 271 | } 272 | 273 | func (input *ForwardInput) spawnDaemon() { 274 | input.logger.Notice("Spawning daemon") 275 | input.wg.Add(1) 276 | go func() { 277 | defer func() { 278 | close(input.shutdownChan) 279 | input.wg.Done() 280 | }() 281 | input.logger.Notice("Daemon started") 282 | loop: 283 | for { 284 | select { 285 | case conn := <-input.acceptChan: 286 | if conn != nil { 287 | input.logger.Notice("Got conn from acceptChan") 288 | newForwardClient(input, input.logger, conn, input.codec).startHandling() 289 | } 290 | case <-input.shutdownChan: 291 | input.listener.Close() 292 | for _, client := range input.clients { 293 | client.shutdown() 294 | } 295 | break loop 296 | } 297 | } 298 | input.logger.Notice("Daemon ended") 299 | }() 300 | } 301 | 302 | func (input *ForwardInput) markCharged(c *forwardClient) { 303 | input.clientsMtx.Lock() 304 | defer input.clientsMtx.Unlock() 305 | input.clients[c.conn] = c 306 | } 307 | 308 | func (input *ForwardInput) markDischarged(c *forwardClient) { 309 | input.clientsMtx.Lock() 310 | defer input.clientsMtx.Unlock() 311 | delete(input.clients, c.conn) 312 | } 313 | 314 | func (input *ForwardInput) String() string { 315 | return "input" 316 | } 317 | 318 | func (input *ForwardInput) Start() { 319 | input.spawnAcceptor() 320 | input.spawnDaemon() 321 | } 322 | 323 | func (input *ForwardInput) WaitForShutdown() { 324 | input.wg.Wait() 325 | } 326 | 327 | func (input *ForwardInput) Stop() { 328 | if atomic.CompareAndSwapUintptr(&input.isShuttingDown, uintptr(0), uintptr(1)) { 329 | input.shutdownChan <- struct{}{} 330 | } 331 | } 332 | 333 | func NewForwardInput(logger *logging.Logger, bind string, port Port) (*ForwardInput, error) { 334 | _codec := codec.MsgpackHandle{} 335 | _codec.MapType = reflect.TypeOf(map[string]interface{}(nil)) 336 | _codec.RawToString = false 337 | addr, err := net.ResolveTCPAddr("tcp", bind) 338 | if err != nil { 339 | logger.Error(err.Error()) 340 | return nil, err 341 | } 342 | listener, err := net.ListenTCP("tcp", addr) 343 | if err != nil { 344 | logger.Error(err.Error()) 345 | return nil, err 346 | } 347 | return &ForwardInput{ 348 | port: port, 349 | logger: logger, 350 | bind: bind, 351 | listener: listener, 352 | codec: &_codec, 353 | clients: make(map[*net.TCPConn]*forwardClient), 354 | clientsMtx: sync.Mutex{}, 355 | entries: 0, 356 | wg: sync.WaitGroup{}, 357 | acceptChan: make(chan *net.TCPConn), 358 | shutdownChan: make(chan struct{}), 359 | isShuttingDown: uintptr(0), 360 | }, nil 361 | } 362 | -------------------------------------------------------------------------------- /output.go: -------------------------------------------------------------------------------- 1 | // 2 | // Fluentd Forwarder 3 | // 4 | // Copyright (C) 2014 Treasure Data, Inc. 5 | // 6 | // Licensed under the Apache License, Version 2.0 (the "License"); 7 | // you may not use this file except in compliance with the License. 8 | // You may obtain a copy of the License at 9 | // 10 | // http://www.apache.org/licenses/LICENSE-2.0 11 | // 12 | // Unless required by applicable law or agreed to in writing, software 13 | // distributed under the License is distributed on an "AS IS" BASIS, 14 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | // See the License for the specific language governing permissions and 16 | // limitations under the License. 17 | // 18 | 19 | package fluentd_forwarder 20 | 21 | import ( 22 | "bytes" 23 | logging "github.com/op/go-logging" 24 | "github.com/ugorji/go/codec" 25 | "io" 26 | "math/rand" 27 | "net" 28 | "os" 29 | "reflect" 30 | "sync" 31 | "sync/atomic" 32 | "time" 33 | ) 34 | 35 | var randSource = rand.NewSource(time.Now().UnixNano()) 36 | 37 | type ForwardOutput struct { 38 | logger *logging.Logger 39 | codec *codec.MsgpackHandle 40 | bind string 41 | retryInterval time.Duration 42 | connectionTimeout time.Duration 43 | writeTimeout time.Duration 44 | enc *codec.Encoder 45 | conn net.Conn 46 | flushInterval time.Duration 47 | wg sync.WaitGroup 48 | journalGroup JournalGroup 49 | journal Journal 50 | emitterChan chan FluentRecordSet 51 | spoolerShutdownChan chan struct{} 52 | isShuttingDown uintptr 53 | completion sync.Cond 54 | hasShutdownCompleted bool 55 | metadata string 56 | } 57 | 58 | func encodeRecordSet(encoder *codec.Encoder, recordSet FluentRecordSet) error { 59 | v := []interface{}{recordSet.Tag, recordSet.Records} 60 | err := encoder.Encode(v) 61 | if err != nil { 62 | return err 63 | } 64 | return err 65 | } 66 | 67 | func (output *ForwardOutput) ensureConnected() error { 68 | if output.conn == nil { 69 | output.logger.Noticef("Connecting to %s...", output.bind) 70 | conn, err := net.DialTimeout("tcp", output.bind, output.connectionTimeout) 71 | if err != nil { 72 | output.logger.Errorf("Failed to connect to %s (reason: %s)", output.bind, err.Error()) 73 | return err 74 | } else { 75 | output.conn = conn 76 | } 77 | } 78 | return nil 79 | } 80 | 81 | func (output *ForwardOutput) sendBuffer(buf []byte) error { 82 | for len(buf) > 0 { 83 | if atomic.LoadUintptr(&output.isShuttingDown) != 0 { 84 | break 85 | } 86 | err := output.ensureConnected() 87 | if err != nil { 88 | output.logger.Infof("Will be retried in %s", output.retryInterval.String()) 89 | time.Sleep(output.retryInterval) 90 | continue 91 | } 92 | startTime := time.Now() 93 | if output.writeTimeout == 0 { 94 | output.conn.SetWriteDeadline(time.Time{}) 95 | } else { 96 | output.conn.SetWriteDeadline(startTime.Add(output.writeTimeout)) 97 | } 98 | n, err := output.conn.Write(buf) 99 | buf = buf[n:] 100 | if err != nil { 101 | output.logger.Errorf("Failed to flush buffer (reason: %s, left: %d bytes)", err.Error(), len(buf)) 102 | err_, ok := err.(net.Error) 103 | if !ok || (!err_.Timeout() && !err_.Temporary()) { 104 | output.conn.Close() 105 | output.conn = nil 106 | continue 107 | } 108 | } 109 | if n > 0 { 110 | elapsed := time.Now().Sub(startTime) 111 | output.logger.Infof("Forwarded %d bytes in %f seconds (%d bytes left)\n", n, elapsed.Seconds(), len(buf)) 112 | } 113 | } 114 | return nil 115 | } 116 | 117 | func (output *ForwardOutput) spawnSpooler() { 118 | output.logger.Notice("Spawning spooler") 119 | output.wg.Add(1) 120 | go func() { 121 | ticker := time.NewTicker(output.flushInterval) 122 | defer func() { 123 | ticker.Stop() 124 | output.journal.Dispose() 125 | if output.conn != nil { 126 | output.conn.Close() 127 | } 128 | output.conn = nil 129 | output.wg.Done() 130 | }() 131 | output.logger.Notice("Spooler started") 132 | outer: 133 | for { 134 | select { 135 | case <-ticker.C: 136 | buf := make([]byte, 16777216) 137 | output.logger.Notice("Flushing...") 138 | err := output.journal.Flush(func(chunk JournalChunk) interface{} { 139 | defer chunk.Dispose() 140 | output.logger.Infof("Flushing chunk %s", chunk.String()) 141 | reader, err := chunk.Reader() 142 | defer reader.Close() 143 | if err != nil { 144 | return err 145 | } 146 | for { 147 | n, err := reader.Read(buf) 148 | if n > 0 { 149 | err_ := output.sendBuffer(buf[:n]) 150 | if err_ != nil { 151 | return err 152 | } 153 | } 154 | if err != nil { 155 | if err == io.EOF { 156 | break 157 | } else { 158 | return err 159 | } 160 | } 161 | } 162 | return nil 163 | }) 164 | if err != nil { 165 | output.logger.Errorf("Error during reading from the journal: %s", err.Error()) 166 | } 167 | case <-output.spoolerShutdownChan: 168 | break outer 169 | } 170 | } 171 | output.logger.Notice("Spooler ended") 172 | }() 173 | } 174 | 175 | func (output *ForwardOutput) spawnEmitter() { 176 | output.logger.Notice("Spawning emitter") 177 | output.wg.Add(1) 178 | go func() { 179 | defer func() { 180 | output.spoolerShutdownChan <- struct{}{} 181 | output.wg.Done() 182 | }() 183 | output.logger.Notice("Emitter started") 184 | buffer := bytes.Buffer{} 185 | for recordSet := range output.emitterChan { 186 | buffer.Reset() 187 | encoder := codec.NewEncoder(&buffer, output.codec) 188 | addMetadata(&recordSet, output.metadata) 189 | err := encodeRecordSet(encoder, recordSet) 190 | if err != nil { 191 | output.logger.Error(err.Error()) 192 | continue 193 | } 194 | output.logger.Debugf("Emitter processed %d entries", len(recordSet.Records)) 195 | output.journal.Write(buffer.Bytes()) 196 | } 197 | output.logger.Notice("Emitter ended") 198 | }() 199 | } 200 | 201 | func (output *ForwardOutput) Emit(recordSets []FluentRecordSet) error { 202 | defer func() { 203 | recover() 204 | }() 205 | for _, recordSet := range recordSets { 206 | output.emitterChan <- recordSet 207 | } 208 | return nil 209 | } 210 | 211 | func (output *ForwardOutput) String() string { 212 | return "output" 213 | } 214 | 215 | func (output *ForwardOutput) Stop() { 216 | if atomic.CompareAndSwapUintptr(&output.isShuttingDown, 0, 1) { 217 | close(output.emitterChan) 218 | } 219 | } 220 | 221 | func (output *ForwardOutput) WaitForShutdown() { 222 | output.completion.L.Lock() 223 | if !output.hasShutdownCompleted { 224 | output.completion.Wait() 225 | } 226 | output.completion.L.Unlock() 227 | } 228 | 229 | func (output *ForwardOutput) Start() { 230 | syncCh := make(chan struct{}) 231 | go func() { 232 | <-syncCh 233 | output.wg.Wait() 234 | err := output.journalGroup.Dispose() 235 | if err != nil { 236 | output.logger.Error(err.Error()) 237 | } 238 | output.completion.L.Lock() 239 | output.hasShutdownCompleted = true 240 | output.completion.Broadcast() 241 | output.completion.L.Unlock() 242 | }() 243 | output.spawnSpooler() 244 | output.spawnEmitter() 245 | syncCh <- struct{}{} 246 | } 247 | 248 | func NewForwardOutput(logger *logging.Logger, bind string, retryInterval time.Duration, connectionTimeout time.Duration, writeTimeout time.Duration, flushInterval time.Duration, journalGroupPath string, maxJournalChunkSize int64, metadata string) (*ForwardOutput, error) { 249 | _codec := codec.MsgpackHandle{} 250 | _codec.MapType = reflect.TypeOf(map[string]interface{}(nil)) 251 | _codec.RawToString = false 252 | _codec.StructToArray = true 253 | 254 | journalFactory := NewFileJournalGroupFactory( 255 | logger, 256 | randSource, 257 | time.Now, 258 | ".log", 259 | os.FileMode(0600), 260 | maxJournalChunkSize, 261 | ) 262 | output := &ForwardOutput{ 263 | logger: logger, 264 | codec: &_codec, 265 | bind: bind, 266 | retryInterval: retryInterval, 267 | connectionTimeout: connectionTimeout, 268 | writeTimeout: writeTimeout, 269 | wg: sync.WaitGroup{}, 270 | flushInterval: flushInterval, 271 | emitterChan: make(chan FluentRecordSet), 272 | spoolerShutdownChan: make(chan struct{}), 273 | isShuttingDown: 0, 274 | completion: sync.Cond{L: &sync.Mutex{}}, 275 | hasShutdownCompleted: false, 276 | metadata: metadata, 277 | } 278 | journalGroup, err := journalFactory.GetJournalGroup(journalGroupPath, output) 279 | if err != nil { 280 | return nil, err 281 | } 282 | output.journalGroup = journalGroup 283 | output.journal = journalGroup.GetJournal("output") 284 | return output, nil 285 | } 286 | -------------------------------------------------------------------------------- /output_td.go: -------------------------------------------------------------------------------- 1 | // 2 | // Fluentd Forwarder 3 | // 4 | // Copyright (C) 2014 Treasure Data, Inc. 5 | // 6 | // Licensed under the Apache License, Version 2.0 (the "License"); 7 | // you may not use this file except in compliance with the License. 8 | // You may obtain a copy of the License at 9 | // 10 | // http://www.apache.org/licenses/LICENSE-2.0 11 | // 12 | // Unless required by applicable law or agreed to in writing, software 13 | // distributed under the License is distributed on an "AS IS" BASIS, 14 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | // See the License for the specific language governing permissions and 16 | // limitations under the License. 17 | // 18 | 19 | package fluentd_forwarder 20 | 21 | import ( 22 | "bytes" 23 | "compress/gzip" 24 | "crypto/x509" 25 | "errors" 26 | ioextras "github.com/moriyoshi/go-ioextras" 27 | logging "github.com/op/go-logging" 28 | td_client "github.com/treasure-data/td-client-go" 29 | "github.com/ugorji/go/codec" 30 | "net" 31 | "os" 32 | "reflect" 33 | "strings" 34 | "sync" 35 | "sync/atomic" 36 | "time" 37 | ) 38 | 39 | type tdOutputSpooler struct { 40 | daemon *tdOutputSpoolerDaemon 41 | ticker *time.Ticker 42 | tag string 43 | databaseName string 44 | tableName string 45 | key string 46 | journal Journal 47 | shutdownChan chan struct{} 48 | isShuttingDown uintptr 49 | client *td_client.TDClient 50 | } 51 | 52 | type tdOutputSpoolerDaemon struct { 53 | output *TDOutput 54 | shutdownChan chan struct{} 55 | spoolersMtx sync.Mutex 56 | spoolers map[string]*tdOutputSpooler 57 | tempFactory ioextras.TempFileRandomAccessStoreFactory 58 | wg sync.WaitGroup 59 | endNotify func(*tdOutputSpoolerDaemon) 60 | } 61 | 62 | type TDOutput struct { 63 | logger *logging.Logger 64 | codec *codec.MsgpackHandle 65 | databaseName string 66 | tableName string 67 | tempDir string 68 | enc *codec.Encoder 69 | conn net.Conn 70 | flushInterval time.Duration 71 | wg sync.WaitGroup 72 | journalGroup JournalGroup 73 | emitterChan chan FluentRecordSet 74 | spoolerDaemon *tdOutputSpoolerDaemon 75 | isShuttingDown uintptr 76 | client *td_client.TDClient 77 | sem chan struct{} 78 | gcChan chan *os.File 79 | completion sync.Cond 80 | hasShutdownCompleted bool 81 | metadata string 82 | } 83 | 84 | func encodeRecords(encoder *codec.Encoder, records []TinyFluentRecord) error { 85 | for _, record := range records { 86 | e := map[string]interface{}{"time": record.Timestamp} 87 | for k, v := range record.Data { 88 | e[k] = v 89 | } 90 | err := encoder.Encode(e) 91 | if err != nil { 92 | return err 93 | } 94 | } 95 | return nil 96 | } 97 | 98 | func (spooler *tdOutputSpooler) cleanup() { 99 | spooler.ticker.Stop() 100 | spooler.journal.Dispose() 101 | spooler.daemon.wg.Done() 102 | } 103 | 104 | func (spooler *tdOutputSpooler) handle() { 105 | defer spooler.cleanup() 106 | spooler.daemon.output.logger.Notice("Spooler started") 107 | outer: 108 | for { 109 | select { 110 | case <-spooler.ticker.C: 111 | spooler.daemon.output.logger.Notice("Flushing...") 112 | err := spooler.journal.Flush(func(chunk JournalChunk) interface{} { 113 | defer chunk.Dispose() 114 | if atomic.LoadUintptr(&spooler.isShuttingDown) != 0 { 115 | return errors.New("Flush aborted") 116 | } 117 | spooler.daemon.output.logger.Infof("Flushing chunk %s", chunk.String()) 118 | size, err := chunk.Size() 119 | if err != nil { 120 | return err 121 | } 122 | if size == 0 { 123 | return nil 124 | } 125 | futureErr := make(chan error, 1) 126 | sem := spooler.daemon.output.sem 127 | sem <- struct{}{} 128 | go func(size int64, chunk JournalChunk, futureErr chan error) { 129 | err := (error)(nil) 130 | defer func() { 131 | if err != nil { 132 | spooler.daemon.output.logger.Infof("Failed to flush chunk %s (reason: %s)", chunk.String(), err.Error()) 133 | } else { 134 | spooler.daemon.output.logger.Infof("Completed flushing chunk %s", chunk.String()) 135 | } 136 | <-sem 137 | // disposal must be done before notifying the initiator 138 | chunk.Dispose() 139 | futureErr <- err 140 | }() 141 | err = func() error { 142 | compressingBlob := NewCompressingBlob( 143 | chunk, 144 | maxInt(4096, int(size/4)), 145 | gzip.BestSpeed, 146 | &spooler.daemon.tempFactory, 147 | ) 148 | defer compressingBlob.Dispose() 149 | _, err := spooler.client.Import( 150 | spooler.databaseName, 151 | spooler.tableName, 152 | "msgpack.gz", 153 | td_client.NewBufferingBlobSize( 154 | compressingBlob, 155 | maxInt(4096, int(size/16)), 156 | ), 157 | chunk.Id(), 158 | ) 159 | return err 160 | }() 161 | }(size, chunk.Dup(), futureErr) 162 | return (<-chan error)(futureErr) 163 | }) 164 | if err != nil { 165 | spooler.daemon.output.logger.Errorf("Error during reading from the journal: %s", err.Error()) 166 | } 167 | case <-spooler.shutdownChan: 168 | break outer 169 | } 170 | } 171 | spooler.daemon.output.logger.Notice("Spooler ended") 172 | } 173 | 174 | func normalizeDatabaseName(name string) (string, error) { 175 | name_ := ([]byte)(name) 176 | if len(name_) == 0 { 177 | return "", errors.New("Empty name is not allowed") 178 | } 179 | if len(name_) < 3 { 180 | name_ = append(name_, ("___"[0 : 3-len(name)])...) 181 | } 182 | if 255 < len(name_) { 183 | name_ = append(name_[0:253], "__"...) 184 | } 185 | name_ = bytes.ToLower(name_) 186 | for i, c := range name_ { 187 | if !((c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '_') { 188 | c = '_' 189 | } 190 | name_[i] = c 191 | } 192 | return (string)(name_), nil 193 | } 194 | 195 | func normalizeTableName(name string) (string, error) { 196 | return normalizeDatabaseName(name) 197 | } 198 | 199 | func newTDOutputSpooler(daemon *tdOutputSpoolerDaemon, databaseName, tableName, key string) *tdOutputSpooler { 200 | journal := daemon.output.journalGroup.GetJournal(key) 201 | return &tdOutputSpooler{ 202 | daemon: daemon, 203 | ticker: time.NewTicker(daemon.output.flushInterval), 204 | databaseName: databaseName, 205 | tableName: tableName, 206 | key: key, 207 | journal: journal, 208 | shutdownChan: make(chan struct{}, 1), 209 | isShuttingDown: 0, 210 | client: daemon.output.client, 211 | } 212 | } 213 | 214 | func (daemon *tdOutputSpoolerDaemon) spawnSpooler(databaseName, tableName, key string) *tdOutputSpooler { 215 | daemon.spoolersMtx.Lock() 216 | defer daemon.spoolersMtx.Unlock() 217 | spooler, exists := daemon.spoolers[key] 218 | if exists { 219 | return spooler 220 | } 221 | spooler = newTDOutputSpooler(daemon, databaseName, tableName, key) 222 | daemon.output.logger.Noticef("Spawning spooler %s", spooler.key) 223 | daemon.spoolers[spooler.key] = spooler 224 | daemon.wg.Add(1) 225 | go spooler.handle() 226 | return spooler 227 | } 228 | 229 | func (daemon *tdOutputSpoolerDaemon) cleanup() { 230 | func() { 231 | daemon.spoolersMtx.Lock() 232 | defer daemon.spoolersMtx.Unlock() 233 | for _, spooler := range daemon.spoolers { 234 | if atomic.CompareAndSwapUintptr(&spooler.isShuttingDown, 0, 1) { 235 | spooler.shutdownChan <- struct{}{} 236 | } 237 | } 238 | }() 239 | daemon.wg.Wait() 240 | if daemon.endNotify != nil { 241 | daemon.endNotify(daemon) 242 | } 243 | daemon.output.wg.Done() 244 | } 245 | 246 | func (daemon *tdOutputSpoolerDaemon) handle() { 247 | defer daemon.cleanup() 248 | daemon.output.logger.Notice("Spooler daemon started") 249 | // spawn Spooler according to the existing journals 250 | for _, key := range daemon.output.journalGroup.GetJournalKeys() { 251 | pair := strings.SplitN(key, ".", 2) 252 | if len(pair) != 2 { 253 | daemon.output.logger.Warningf("Journal %s ignored", key) 254 | continue 255 | } 256 | daemon.spawnSpooler(pair[0], pair[1], key) 257 | } 258 | outer: 259 | for { 260 | select { 261 | case <-daemon.shutdownChan: 262 | break outer 263 | } 264 | } 265 | daemon.output.logger.Notice("Spooler daemon ended") 266 | } 267 | 268 | func newTDOutputSpoolerDaemon(output *TDOutput) *tdOutputSpoolerDaemon { 269 | return &tdOutputSpoolerDaemon{ 270 | output: output, 271 | shutdownChan: make(chan struct{}, 1), 272 | spoolers: make(map[string]*tdOutputSpooler), 273 | tempFactory: ioextras.TempFileRandomAccessStoreFactory{output.tempDir, "", output.gcChan}, 274 | wg: sync.WaitGroup{}, 275 | endNotify: func(*tdOutputSpoolerDaemon) { 276 | close(output.gcChan) 277 | }, 278 | } 279 | } 280 | 281 | func (output *TDOutput) spawnSpoolerDaemon() { 282 | output.logger.Notice("Spawning spooler daemon") 283 | output.spoolerDaemon = newTDOutputSpoolerDaemon(output) 284 | output.wg.Add(1) 285 | go output.spoolerDaemon.handle() 286 | } 287 | 288 | func (daemon *tdOutputSpoolerDaemon) getSpooler(tag string) (*tdOutputSpooler, error) { 289 | databaseName := daemon.output.databaseName 290 | tableName := daemon.output.tableName 291 | if databaseName == "*" { 292 | if tableName == "*" { 293 | c := strings.SplitN(tag, ".", 2) 294 | if len(c) == 1 { 295 | tableName = c[0] 296 | } else if len(c) == 2 { 297 | databaseName = c[0] 298 | tableName = c[1] 299 | } 300 | } else { 301 | databaseName = tag 302 | } 303 | } else { 304 | if tableName == "*" { 305 | tableName = tag 306 | } 307 | } 308 | databaseName, err := normalizeDatabaseName(databaseName) 309 | if err != nil { 310 | return nil, err 311 | } 312 | tableName, err = normalizeTableName(tableName) 313 | if err != nil { 314 | return nil, err 315 | } 316 | key := databaseName + "." + tableName 317 | return daemon.spawnSpooler(databaseName, tableName, key), nil 318 | } 319 | 320 | func (output *TDOutput) spawnEmitter() { 321 | output.logger.Notice("Spawning emitter") 322 | output.wg.Add(1) 323 | go func() { 324 | defer func() { 325 | output.spoolerDaemon.shutdownChan <- struct{}{} 326 | output.wg.Done() 327 | }() 328 | output.logger.Notice("Emitter started") 329 | buffer := bytes.Buffer{} 330 | for recordSet := range output.emitterChan { 331 | buffer.Reset() 332 | encoder := codec.NewEncoder(&buffer, output.codec) 333 | err := func() error { 334 | spooler, err := output.spoolerDaemon.getSpooler(recordSet.Tag) 335 | if err != nil { 336 | return err 337 | } 338 | addMetadata(&recordSet, output.metadata) 339 | err = encodeRecords(encoder, recordSet.Records) 340 | if err != nil { 341 | return err 342 | } 343 | output.logger.Debugf("Emitter processed %d entries", len(recordSet.Records)) 344 | return spooler.journal.Write(buffer.Bytes()) 345 | }() 346 | if err != nil { 347 | output.logger.Error(err.Error()) 348 | continue 349 | } 350 | } 351 | output.logger.Notice("Emitter ended") 352 | }() 353 | } 354 | 355 | func (output *TDOutput) spawnTempFileCollector() { 356 | output.logger.Notice("Spawning temporary file collector") 357 | output.wg.Add(1) 358 | go func() { 359 | defer func() { 360 | output.wg.Done() 361 | }() 362 | for f := range output.gcChan { 363 | output.logger.Debugf("Deleting %s...", f.Name()) 364 | err := os.Remove(f.Name()) 365 | if err != nil { 366 | output.logger.Warningf("Failed to delete %s: %s", f.Name(), err.Error()) 367 | } 368 | } 369 | output.logger.Debug("temporary file collector ended") 370 | }() 371 | } 372 | 373 | func (output *TDOutput) Emit(recordSets []FluentRecordSet) error { 374 | defer func() { 375 | recover() 376 | }() 377 | for _, recordSet := range recordSets { 378 | output.emitterChan <- recordSet 379 | } 380 | return nil 381 | } 382 | 383 | func (output *TDOutput) String() string { 384 | return "output" 385 | } 386 | 387 | func (output *TDOutput) Stop() { 388 | if atomic.CompareAndSwapUintptr(&output.isShuttingDown, 0, 1) { 389 | close(output.emitterChan) 390 | } 391 | } 392 | 393 | func (output *TDOutput) WaitForShutdown() { 394 | output.completion.L.Lock() 395 | if !output.hasShutdownCompleted { 396 | output.completion.Wait() 397 | } 398 | output.completion.L.Unlock() 399 | } 400 | 401 | func (output *TDOutput) Start() { 402 | syncCh := make(chan struct{}) 403 | go func() { 404 | <-syncCh 405 | output.wg.Wait() 406 | err := output.journalGroup.Dispose() 407 | if err != nil { 408 | output.logger.Error(err.Error()) 409 | } 410 | output.completion.L.Lock() 411 | output.hasShutdownCompleted = true 412 | output.completion.Broadcast() 413 | output.completion.L.Unlock() 414 | }() 415 | output.spawnTempFileCollector() 416 | output.spawnEmitter() 417 | output.spawnSpoolerDaemon() 418 | syncCh <- struct{}{} 419 | } 420 | 421 | func NewTDOutput( 422 | logger *logging.Logger, 423 | endpoint string, 424 | connectionTimeout time.Duration, 425 | writeTimeout time.Duration, 426 | flushInterval time.Duration, 427 | parallelism int, 428 | journalGroupPath string, 429 | maxJournalChunkSize int64, 430 | apiKey string, 431 | databaseName string, 432 | tableName string, 433 | tempDir string, 434 | useSsl bool, 435 | rootCAs *x509.CertPool, 436 | httpProxy string, 437 | metadata string, 438 | ) (*TDOutput, error) { 439 | _codec := codec.MsgpackHandle{} 440 | _codec.MapType = reflect.TypeOf(map[string]interface{}(nil)) 441 | _codec.RawToString = false 442 | _codec.StructToArray = true 443 | 444 | journalFactory := NewFileJournalGroupFactory( 445 | logger, 446 | randSource, 447 | time.Now, 448 | ".log", 449 | os.FileMode(0600), 450 | maxJournalChunkSize, 451 | ) 452 | router := (td_client.EndpointRouter)(nil) 453 | if endpoint != "" { 454 | router = &td_client.FixedEndpointRouter{endpoint} 455 | } 456 | httpProxy_ := (interface{})(nil) 457 | if httpProxy != "" { 458 | httpProxy_ = httpProxy 459 | } 460 | client, err := td_client.NewTDClient(td_client.Settings{ 461 | ApiKey: apiKey, 462 | Router: router, 463 | ConnectionTimeout: connectionTimeout, 464 | // ReadTimeout: readTimeout, // TODO 465 | SendTimeout: writeTimeout, 466 | Ssl: useSsl, 467 | RootCAs: rootCAs, 468 | Proxy: httpProxy_, 469 | }) 470 | if err != nil { 471 | return nil, err 472 | } 473 | output := &TDOutput{ 474 | logger: logger, 475 | codec: &_codec, 476 | wg: sync.WaitGroup{}, 477 | flushInterval: flushInterval, 478 | emitterChan: make(chan FluentRecordSet), 479 | isShuttingDown: 0, 480 | client: client, 481 | databaseName: databaseName, 482 | tableName: tableName, 483 | tempDir: tempDir, 484 | sem: make(chan struct{}, parallelism), 485 | gcChan: make(chan *os.File, 10), 486 | completion: sync.Cond{L: &sync.Mutex{}}, 487 | hasShutdownCompleted: false, 488 | metadata: metadata, 489 | } 490 | journalGroup, err := journalFactory.GetJournalGroup(journalGroupPath, output) 491 | if err != nil { 492 | return nil, err 493 | } 494 | output.journalGroup = journalGroup 495 | return output, nil 496 | } 497 | -------------------------------------------------------------------------------- /output_td_test.go: -------------------------------------------------------------------------------- 1 | // 2 | // Fluentd Forwarder 3 | // 4 | // Copyright (C) 2014 Treasure Data, Inc. 5 | // 6 | // Licensed under the Apache License, Version 2.0 (the "License"); 7 | // you may not use this file except in compliance with the License. 8 | // You may obtain a copy of the License at 9 | // 10 | // http://www.apache.org/licenses/LICENSE-2.0 11 | // 12 | // Unless required by applicable law or agreed to in writing, software 13 | // distributed under the License is distributed on an "AS IS" BASIS, 14 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | // See the License for the specific language governing permissions and 16 | // limitations under the License. 17 | // 18 | 19 | package fluentd_forwarder 20 | 21 | import "testing" 22 | 23 | func TestNormalizeDatabaseName(t *testing.T) { 24 | { 25 | _, err := normalizeDatabaseName("") 26 | if err == nil { 27 | t.FailNow() 28 | } 29 | } 30 | { 31 | result, err := normalizeDatabaseName("a") 32 | if err != nil { 33 | t.Logf("%s", err.Error()) 34 | t.FailNow() 35 | } 36 | t.Logf("%s", result) 37 | if result != "a__" { 38 | t.Fail() 39 | } 40 | } 41 | { 42 | result, err := normalizeDatabaseName("ab") 43 | if err != nil { 44 | t.Logf("%s", err.Error()) 45 | t.FailNow() 46 | } 47 | t.Logf("%s", result) 48 | if result != "ab_" { 49 | t.Fail() 50 | } 51 | } 52 | { 53 | result, err := normalizeDatabaseName("abc") 54 | if err != nil { 55 | t.Logf("%s", err.Error()) 56 | t.FailNow() 57 | } 58 | t.Logf("%s", result) 59 | if result != "abc" { 60 | t.Fail() 61 | } 62 | } 63 | { 64 | result, err := normalizeDatabaseName("abcd") 65 | if err != nil { 66 | t.Logf("%s", err.Error()) 67 | t.FailNow() 68 | } 69 | t.Logf("%s", result) 70 | if result != "abcd" { 71 | t.Fail() 72 | } 73 | } 74 | { 75 | result, err := normalizeDatabaseName("abc.def.ghi") 76 | if err != nil { 77 | t.Logf("%s", err.Error()) 78 | t.FailNow() 79 | } 80 | t.Logf("%s", result) 81 | if result != "abc_def_ghi" { 82 | t.Fail() 83 | } 84 | } 85 | } 86 | -------------------------------------------------------------------------------- /path_builder.go: -------------------------------------------------------------------------------- 1 | // 2 | // Copyright (C) 2014 Moriyoshi Koizumi 3 | // Copyright (C) 2014 Treasure Data, Inc. 4 | // 5 | // Permission is hereby granted, free of charge, to any person obtaining a copy 6 | // of this software and associated documentation files (the "Software"), to deal 7 | // in the Software without restriction, including without limitation the rights 8 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | // copies of the Software, and to permit persons to whom the Software is 10 | // furnished to do so, subject to the following conditions: 11 | // 12 | // The above copyright notice and this permission notice shall be included in 13 | // all copies or substantial portions of the Software. 14 | // 15 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 21 | // THE SOFTWARE. 22 | // 23 | 24 | package fluentd_forwarder 25 | 26 | import ( 27 | "errors" 28 | "fmt" 29 | "net/url" 30 | "regexp" 31 | "strconv" 32 | "strings" 33 | "time" 34 | ) 35 | 36 | type JournalFileType rune 37 | 38 | const ( 39 | Head = JournalFileType('b') 40 | Rest = JournalFileType('q') 41 | ) 42 | 43 | type JournalPathInfo struct { 44 | Key string 45 | Type JournalFileType 46 | VariablePortion string 47 | TSuffix string 48 | Timestamp int64 // elapsed time in msec since epoch 49 | UniqueId []byte 50 | } 51 | 52 | var NilJournalPathInfo = JournalPathInfo{"", 0, "", "", 0, nil} 53 | 54 | var pathRegexp, _ = regexp.Compile(fmt.Sprintf("^(.*)[._](%c|%c)([0-9a-fA-F]{1,32})$", Head, Rest)) 55 | 56 | func encodeKey(key string) string { 57 | keyLen := len(key) 58 | retval := make([]byte, keyLen*2) 59 | i := 0 60 | for j := 0; j < keyLen; j += 1 { 61 | c := key[j] 62 | if c == '-' || c == '_' || c == '.' || (c >= '0' && c <= '9') || (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') { 63 | cap_ := cap(retval) 64 | if i >= cap_ { 65 | newRetval := make([]byte, cap_+len(key)) 66 | copy(newRetval, retval) 67 | retval = newRetval 68 | } 69 | retval[i] = c 70 | i += 1 71 | } else { 72 | cap_ := cap(retval) 73 | if i+3 > cap_ { 74 | newSize := cap_ 75 | for { 76 | newSize += len(key) 77 | if i+3 <= newSize { 78 | break 79 | } 80 | if newSize < cap_ { 81 | // unlikely 82 | panic("?") 83 | } 84 | } 85 | newRetval := make([]byte, newSize) 86 | copy(newRetval, retval) 87 | retval = newRetval 88 | } 89 | retval[i] = '%' 90 | retval[i+1] = "0123456789abcdef"[(c>>4)&15] 91 | retval[i+2] = "0123456789abcdef"[c&15] 92 | i += 3 93 | } 94 | } 95 | return string(retval[0:i]) 96 | } 97 | 98 | func decodeKey(encoded string) (string, error) { 99 | return url.QueryUnescape(encoded) 100 | } 101 | 102 | func convertTSuffixToUniqueId(tSuffix string) ([]byte, error) { 103 | tSuffixLen := len(tSuffix) 104 | buf := make([]byte, tSuffixLen) 105 | for i := 0; i < tSuffixLen; i += 2 { 106 | ii, err := strconv.ParseUint(tSuffix[i:i+2], 16, 8) 107 | if err != nil { 108 | return nil, err 109 | } 110 | buf[i/2] = byte(ii) 111 | buf[(i+tSuffixLen)/2] = byte(ii) 112 | } 113 | return buf, nil 114 | } 115 | 116 | func convertTSuffixToUnixNano(tSuffix string) (int64, error) { 117 | t, err := strconv.ParseInt(tSuffix, 16, 64) 118 | return t >> 12, err 119 | } 120 | 121 | func IsValidJournalPathInfo(info JournalPathInfo) bool { 122 | return len(info.Key) > 0 && info.Type != 0 123 | } 124 | 125 | func BuildJournalPathWithTSuffix(key string, bq JournalFileType, tSuffix string) string { 126 | encodedKey := encodeKey(key) 127 | return fmt.Sprintf( 128 | "%s.%c%s", 129 | encodedKey, 130 | rune(bq), 131 | tSuffix, 132 | ) 133 | } 134 | 135 | func BuildJournalPath(key string, bq JournalFileType, time_ time.Time, randValue int64) JournalPathInfo { 136 | timestamp := time_.UnixNano() 137 | t := (timestamp/1000)<<12 | (randValue & 0xfff) 138 | tSuffix := strconv.FormatInt(t, 16) 139 | if pad := 16 - len(tSuffix); pad > 0 { 140 | // unlikely 141 | tSuffix = strings.Repeat("0", pad) + tSuffix 142 | } 143 | uniqueId, err := convertTSuffixToUniqueId(tSuffix) 144 | if err != nil { 145 | panic("WTF? " + err.Error()) 146 | } // should never happen 147 | return JournalPathInfo{ 148 | Key: key, 149 | Type: bq, 150 | VariablePortion: BuildJournalPathWithTSuffix(key, bq, tSuffix), 151 | TSuffix: tSuffix, 152 | Timestamp: timestamp, 153 | UniqueId: uniqueId, 154 | } 155 | } 156 | 157 | func firstRune(s string) rune { 158 | for _, r := range s { 159 | return r 160 | } 161 | return -1 // in case the string is empty 162 | } 163 | 164 | func DecodeJournalPath(variablePortion string) (JournalPathInfo, error) { 165 | m := pathRegexp.FindStringSubmatch(variablePortion) 166 | if m == nil { 167 | return NilJournalPathInfo, errors.New("malformed path string") 168 | } 169 | key, err := decodeKey(m[1]) 170 | if err != nil { 171 | return NilJournalPathInfo, errors.New("malformed path string") 172 | } 173 | uniqueId, err := convertTSuffixToUniqueId(m[3]) 174 | if err != nil { 175 | return NilJournalPathInfo, errors.New("malformed path string") 176 | } 177 | timestamp, err := convertTSuffixToUnixNano(m[3]) 178 | if err != nil { 179 | return NilJournalPathInfo, errors.New("malformed path string") 180 | } 181 | return JournalPathInfo{ 182 | Key: key, 183 | Type: JournalFileType(firstRune(m[2])), 184 | VariablePortion: variablePortion, 185 | TSuffix: m[3], 186 | Timestamp: timestamp, 187 | UniqueId: uniqueId, 188 | }, nil 189 | } 190 | -------------------------------------------------------------------------------- /path_builder_test.go: -------------------------------------------------------------------------------- 1 | // 2 | // Copyright (C) 2014 Moriyoshi Koizumi 3 | // Copyright (C) 2014 Treasure Data, Inc. 4 | // 5 | // Permission is hereby granted, free of charge, to any person obtaining a copy 6 | // of this software and associated documentation files (the "Software"), to deal 7 | // in the Software without restriction, including without limitation the rights 8 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | // copies of the Software, and to permit persons to whom the Software is 10 | // furnished to do so, subject to the following conditions: 11 | // 12 | // The above copyright notice and this permission notice shall be included in 13 | // all copies or substantial portions of the Software. 14 | // 15 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 21 | // THE SOFTWARE. 22 | // 23 | 24 | package fluentd_forwarder 25 | 26 | import ( 27 | "testing" 28 | "time" 29 | ) 30 | 31 | func Test_BuildJournalPath(t *testing.T) { 32 | info := BuildJournalPath( 33 | "test", 34 | JournalFileType('b'), 35 | time.Date(2014, 1, 1, 0, 0, 0, 0, time.UTC), 36 | 0x0, 37 | ) 38 | t.Logf("%+v", info) 39 | if len(info.UniqueId) != len(info.TSuffix) { 40 | t.Fail() 41 | } 42 | if info.VariablePortion != "test.b4eedd5baba000000" { 43 | t.Fail() 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /utils.go: -------------------------------------------------------------------------------- 1 | // 2 | // Fluentd Forwarder 3 | // 4 | // Copyright (C) 2014 Treasure Data, Inc. 5 | // 6 | // Licensed under the Apache License, Version 2.0 (the "License"); 7 | // you may not use this file except in compliance with the License. 8 | // You may obtain a copy of the License at 9 | // 10 | // http://www.apache.org/licenses/LICENSE-2.0 11 | // 12 | // Unless required by applicable law or agreed to in writing, software 13 | // distributed under the License is distributed on an "AS IS" BASIS, 14 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | // See the License for the specific language governing permissions and 16 | // limitations under the License. 17 | // 18 | 19 | package fluentd_forwarder 20 | 21 | func maxInt(a, b int) int { 22 | if a >= b { 23 | return a 24 | } else { 25 | return b 26 | } 27 | } 28 | 29 | func addMetadata(recordSet *FluentRecordSet, metadata string) { 30 | if metadata != "" { 31 | for _, record := range recordSet.Records { 32 | record.Data["metadata"] = string(metadata) 33 | } 34 | } 35 | } 36 | -------------------------------------------------------------------------------- /worker_set.go: -------------------------------------------------------------------------------- 1 | // 2 | // Fluentd Forwarder 3 | // 4 | // Copyright (C) 2014 Treasure Data, Inc. 5 | // 6 | // Licensed under the Apache License, Version 2.0 (the "License"); 7 | // you may not use this file except in compliance with the License. 8 | // You may obtain a copy of the License at 9 | // 10 | // http://www.apache.org/licenses/LICENSE-2.0 11 | // 12 | // Unless required by applicable law or agreed to in writing, software 13 | // distributed under the License is distributed on an "AS IS" BASIS, 14 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | // See the License for the specific language governing permissions and 16 | // limitations under the License. 17 | // 18 | 19 | package fluentd_forwarder 20 | 21 | import ( 22 | "sync" 23 | ) 24 | 25 | type WorkerSet struct { 26 | workers map[Worker]struct{} 27 | mtx sync.Mutex 28 | } 29 | 30 | func (set *WorkerSet) Add(worker Worker) { 31 | set.mtx.Lock() 32 | defer set.mtx.Unlock() 33 | set.workers[worker] = struct{}{} 34 | } 35 | 36 | func (set *WorkerSet) Remove(worker Worker) { 37 | set.mtx.Lock() 38 | defer set.mtx.Unlock() 39 | delete(set.workers, worker) 40 | } 41 | 42 | func (set *WorkerSet) Slice() []Worker { 43 | retval := make([]Worker, len(set.workers)) 44 | set.mtx.Lock() 45 | defer set.mtx.Unlock() 46 | i := 0 47 | for worker, _ := range set.workers { 48 | retval[i] = worker 49 | i += 1 50 | } 51 | return retval 52 | } 53 | 54 | func NewWorkerSet() *WorkerSet { 55 | return &WorkerSet{ 56 | workers: make(map[Worker]struct{}), 57 | mtx: sync.Mutex{}, 58 | } 59 | } 60 | --------------------------------------------------------------------------------