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 = &timestamp.Timestamp{
   165  		Nanos:   int32((milliSeconds % 1000) * 1000000),
   166  		Seconds: milliSeconds / 1000,
   167  	}
   168  
   169  	return logsprayMsg, nil
   170  }