jkawamoto/roadie

//
// cloud/azure/log_manager.go
//
// Copyright (c) 2016-2017 Junpei Kawamoto
//
// This file is part of Roadie.
//
// Roadie is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Roadie is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with Roadie.  If not, see <http://www.gnu.org/licenses/>.
//

package azure

import (
    "bufio"
    "context"
    "fmt"
    "io"
    "log"
    "net/url"
    "path"
    "strings"
    "time"

    "golang.org/x/sync/errgroup"

    "github.com/Azure/azure-sdk-for-go/storage"
    "github.com/jkawamoto/roadie/cloud"
    "github.com/jkawamoto/roadie/script"
)

// LogManager defines a service interface for obtaining log entries.
type LogManager struct {
    storage *StorageService
    batch   *BatchService
    Config  *Config
    Logger  *log.Logger
}

// NewLogManager creates a new log manger for Azure.
func NewLogManager(ctx context.Context, cfg *Config, logger *log.Logger) (m *LogManager, err error) {

    storage, err := NewStorageService(ctx, cfg, logger)
    if err != nil {
        return
    }
    batch, err := NewBatchService(ctx, cfg, logger)
    if err != nil {
        return
    }

    m = &LogManager{
        storage: storage,
        batch:   batch,
        Config:  cfg,
        Logger:  logger,
    }
    return

}

// Get retrievs log entries.
func (m *LogManager) Get(ctx context.Context, instanceName string, from time.Time, handler cloud.LogHandler) (err error) {

    var urls []*url.URL
    var loc *url.URL
    for _, format := range []string{"%v-init.log", "%v.log"} {
        loc, err = url.Parse(script.RoadieSchemePrefix + path.Join(LogContainer, fmt.Sprintf(format, instanceName)))
        if err != nil {
            return
        }
        urls = append(urls, loc)
    }
    return m.get(ctx, urls, handler)

}

// get retrieves log files represented by a given URLs and sends each line to a given handler.
func (m *LogManager) get(ctx context.Context, urls []*url.URL, handler cloud.LogHandler) (err error) {

    ch := make(chan string)
    wg, ctx := errgroup.WithContext(ctx)
    reader, writer := io.Pipe()

    wg.Go(func() error {
        defer reader.Close()
        scanner := bufio.NewScanner(reader)
        for scanner.Scan() {
            ch <- scanner.Text()
        }
        close(ch)
        return nil
    })

    wg.Go(func() (err error) {
        var ignore bool
        for {
            select {
            case <-ctx.Done():
                reader.Close()
                writer.Close()
                ignore = true
                err = ctx.Err()
            case line, ok := <-ch:
                if !ok {
                    return
                }
                if !ignore {
                    fields := strings.SplitN(line, " ", 3)
                    if len(fields) != 3 {
                        continue
                    }
                    t, _ := time.Parse("2006/01/02 15:04:05", fmt.Sprintf("%v %v", fields[0], fields[1]))
                    err = handler(t.UTC().In(time.Local), fields[2], false)
                    if err != nil {
                        reader.Close()
                        writer.Close()
                        ignore = true
                    }
                }
            }
        }
    })

    wg.Go(func() (err error) {
        defer writer.Close()
        for _, loc := range urls {
            err = ignoreNotFoundError(m.storage.Download(ctx, loc, writer))
            if err != nil {
                break
            }
        }
        return
    })

    return wg.Wait()

}

// Delete instance log.
func (m *LogManager) Delete(ctx context.Context, instanceName string) (err error) {

    var loc *url.URL
    // Delete a config file for initialization.
    loc, err = url.Parse(script.RoadieSchemePrefix + StartupContainer)
    if err != nil {
        return
    }
    err = m.storage.List(ctx, loc, func(info *cloud.FileInfo) error {
        if strings.HasPrefix(info.Name, instanceName) && strings.HasSuffix(info.Name, "-init.cfg") {
            return m.storage.Delete(ctx, info.URL)
        }
        return nil
    })
    if err != nil {
        return
    }

    // Delete log files.
    loc, err = url.Parse(script.RoadieSchemePrefix + LogContainer)
    if err != nil {
        return
    }
    err = m.storage.List(ctx, loc, func(info *cloud.FileInfo) error {
        if strings.HasPrefix(info.Name, instanceName) || strings.HasPrefix(info.Name, fmt.Sprintf("task-%v", instanceName)) {
            return m.storage.Delete(ctx, info.URL)
        }
        return nil
    })
    if err != nil {
        return
    }

    // Delete the job.
    jobSet, err := m.batch.Jobs(ctx)
    if err != nil {
        return
    }
    for name := range jobSet {
        if name == instanceName {
            err = m.batch.DeleteJob(ctx, instanceName)
            if err != nil {
                return
            }
        }
    }
    return nil

}

// GetQueueLog retrieves log of a given queue.
func (m *LogManager) GetQueueLog(ctx context.Context, queue string, handler cloud.LogHandler) (err error) {

    queue = queueName(queue)
    tasks, err := m.batch.Tasks(ctx, queue)
    if err != nil {
        return
    }

    var urls []*url.URL
    var loc *url.URL
    loc, err = url.Parse(script.RoadieSchemePrefix + path.Join(LogContainer, fmt.Sprintf("%v-init.log", queue)))
    if err != nil {
        return
    }
    urls = append(urls, loc)

    for _, task := range tasks {
        loc, err = url.Parse(script.RoadieSchemePrefix + path.Join(LogContainer, fmt.Sprintf("%v.log", task.ID)))
        if err != nil {
            return
        }
        urls = append(urls, loc)
    }

    return m.get(ctx, urls, handler)

}

// GetTaskLog retrieves log of a given task.
func (m *LogManager) GetTaskLog(ctx context.Context, queue, task string, handler cloud.LogHandler) (err error) {

    // queue = queueName(queue)
    task = taskName(task)
    loc, err := url.Parse(script.RoadieSchemePrefix + path.Join(LogContainer, fmt.Sprintf("%v.log", task)))
    if err != nil {
        return
    }
    return m.get(ctx, []*url.URL{loc}, handler)

}

// ignoreNotFoundError is a wrapper function and ignores not found error.
func ignoreNotFoundError(err error) error {

    if err != nil {
        switch e := err.(type) {
        case storage.AzureStorageServiceError:
            if e.StatusCode == 404 {
                err = nil
            }
        }
    }
    return err

}