github.com/qubitproducts/logspray@v0.2.14/sources/kinesis/kinesissource.go (about) 1 // Copyright 2016 Qubit Digital Ltd. 2 // Licensed under the Apache License, Version 2.0 (the "License"); 3 // you may not use this file except in compliance with the License. 4 // You may obtain a copy of the License at 5 // 6 // http://www.apache.org/licenses/LICENSE-2.0 7 // 8 // Unless required by applicable law or agreed to in writing, software 9 // distributed under the License is distributed on an "AS IS" BASIS, 10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 11 // See the License for the specific language governing permissions and 12 // limitations under the License. 13 // Package logspray is a collection of tools for streaming and indexing 14 // large volumes of dynamic logs. 15 16 package kinesis 17 18 import ( 19 "bytes" 20 "compress/gzip" 21 "context" 22 "encoding/json" 23 24 "github.com/QubitProducts/logspray/proto/logspray" 25 "github.com/QubitProducts/logspray/sources" 26 "github.com/aws/aws-sdk-go/aws" 27 "github.com/aws/aws-sdk-go/service/kinesis" 28 "github.com/golang/glog" 29 "github.com/golang/protobuf/ptypes/timestamp" 30 "github.com/pkg/errors" 31 ) 32 33 // MessageReader is a log source that reads from kinesis shard. 34 type MessageReader struct { 35 shardIterator *string 36 kinesis *kinesis.Kinesis 37 stream string 38 messagesChannel chan Message 39 } 40 41 type Message struct { 42 logsprayMsg *logspray.Message 43 error error 44 } 45 46 type KinesisMessage struct { 47 LogEvents []interface{} `json:"logEvents"` 48 } 49 50 // ReadTarget creates a new log source from a kinesis shard 51 func (w *Watcher) ReadTarget(ctx context.Context, shardId string, fromStart bool) (sources.MessageReader, error) { 52 shardIterator, err := w.shardIterator(ctx, shardId, fromStart) 53 if err != nil { 54 return nil, err 55 } 56 messagesChannel := make(chan Message, w.messagesChannelSize) 57 58 msgReader := &MessageReader{ 59 shardIterator: shardIterator, 60 kinesis: w.kinesis, 61 stream: w.stream, 62 messagesChannel: messagesChannel, 63 } 64 65 go msgReader.startReadingFromKinesis(ctx) 66 67 return msgReader, err 68 } 69 70 // MessageRead implements the LogSourcer interface 71 func (mr *MessageReader) MessageRead(ctx context.Context) (*logspray.Message, error) { 72 message := <-mr.messagesChannel 73 return message.logsprayMsg, message.error 74 } 75 76 func (w *Watcher) shardIterator(ctx context.Context, shardId string, fromStart bool) (*string, error) { 77 shardIteratorType := aws.String("LATEST") 78 if fromStart { 79 shardIteratorType = aws.String("TRIM_HORIZON") 80 } 81 82 resp, err := w.kinesis.GetShardIteratorWithContext(ctx, &kinesis.GetShardIteratorInput{ 83 ShardId: aws.String(shardId), 84 ShardIteratorType: shardIteratorType, 85 StreamName: aws.String(w.stream), 86 }) 87 if err != nil { 88 return nil, errors.Wrap(err, "could not get shard iterator") 89 } 90 91 return resp.ShardIterator, ctx.Err() 92 } 93 94 func (mr *MessageReader) startReadingFromKinesis(ctx context.Context) { 95 defer close(mr.messagesChannel) 96 for { 97 resp, kinesisErr := mr.kinesis.GetRecordsWithContext(ctx, &kinesis.GetRecordsInput{ 98 ShardIterator: mr.shardIterator, 99 Limit: aws.Int64(64), 100 }) 101 102 if kinesisErr != nil { 103 mr.messagesChannel <- Message{nil, kinesisErr} 104 continue 105 } 106 107 mr.shardIterator = resp.NextShardIterator 108 109 for _, r := range resp.Records { 110 // decompressing data from kinesis 111 reader := bytes.NewReader(r.Data) 112 gzipReader, gzipErr := gzip.NewReader(reader) 113 if gzipErr != nil { 114 glog.Error(gzipErr) 115 continue 116 } 117 118 kinesisMsg := KinesisMessage{} 119 buf := new(bytes.Buffer) 120 buf.ReadFrom(gzipReader) 121 marshlingErr := json.Unmarshal(buf.Bytes(), &kinesisMsg) 122 gzipReader.Close() 123 if marshlingErr != nil { 124 glog.Error(marshlingErr) 125 continue 126 } 127 128 for _, log := range kinesisMsg.LogEvents { 129 logsprayMsg, parseErr := parseLog(log) 130 131 if parseErr != nil { 132 glog.Error(parseErr) 133 continue 134 } 135 136 msg := Message{logsprayMsg, nil} 137 138 select { 139 case <-ctx.Done(): 140 return 141 case mr.messagesChannel <- msg: 142 } 143 } 144 } 145 } 146 } 147 148 func parseLog(log interface{}) (*logspray.Message, error) { 149 logEvents, ok := log.(map[string]interface{}) 150 if !ok { 151 return nil, errors.New("Failed to parse kinesis message") 152 } 153 154 message, ok := logEvents["message"].(string) 155 if !ok { 156 return nil, errors.New("Failed to parse message field in kinesis message") 157 } 158 159 milliSeconds := int64(logEvents["timestamp"].(float64)) 160 161 logsprayMsg := &logspray.Message{} 162 logsprayMsg.Text = message 163 164 logsprayMsg.Time = ×tamp.Timestamp{ 165 Nanos: int32((milliSeconds % 1000) * 1000000), 166 Seconds: milliSeconds / 1000, 167 } 168 169 return logsprayMsg, nil 170 }