2
\$\begingroup\$

EDIT: I have updated the code, if you have comments about the previous version of the code, please do post them anyway.

The task is to determine the country by an IP.

Provided programme queries two external services (namely api.ipstack.com and geoip.nekudo.com) to determine the country by IP. As services has their limits of maximum queries per minute, the programme switches between the services during the execution.

Any suggestion and critique would be very welcome!

package cquery

import (
    "encoding/json"
    "fmt"
    "io"
    "net/http"
    "net/url"
    "os"
    "runtime"
    "sync"
    "time"

    "github.com/patrickmn/go-cache"
    "github.com/sirupsen/logrus"
)

// client - core type of the package
type client struct {
    *Config
    *logrus.Logger
    *http.Client
    *cache.Cache
    workersPull chan struct{}
}

// Config - configuration of the client
type Config = struct {
    CacheExpiration time.Duration
    LogLevel        logrus.Level
    Timeout         time.Duration // waiting time for a response from the service
    MaxWorkers      int           // maximum number of simultaneously running workers
    Services
}

// Service - describes parameters of a Service which provides information about the country by the provided IP
type Service struct {
    Key         string
    Host        string
    Path        string
    Scheme      string
    QueryPeriod time.Duration
    QueryLimit  int64     // maximum amount of queries per QueryPeriod
    QueryLap    time.Time // timestamp of the first query
    QueryRemain int64     // number of remaining queries for the given period
    mux         sync.Mutex
    Wait        time.Duration
}

type Services map[string]*Service

// New - creates new instance of client type
func New(cfg *Config) *client {
    log := logrus.New()
    logFile, _ := os.Create("cquery.log")
    log.SetLevel(cfg.LogLevel)
    mw := io.MultiWriter(os.Stdout, logFile)
    log.Out = mw
    for _, service := range cfg.Services {
        service.QueryRemain = service.QueryLimit
    }
    return &client{
        Config:      cfg,
        Logger:      log,
        Client:      &http.Client{},
        Cache:       cache.New(cfg.CacheExpiration*time.Second, cfg.CacheExpiration*time.Second),
        workersPull: make(chan struct{}, cfg.MaxWorkers-1), // set the pul of workers
    }
}

// selectService - selects a service for query
func (c *client) selectService() (result string) {

    for service, config := range c.Services {
        // if the first query has been made too long ago, nullify the statistics
        if time.Since(config.QueryLap) > config.QueryPeriod {
            config.mux.Lock()
            config.QueryRemain = config.QueryLimit
            config.mux.Unlock()
        }
        // return any service with unexhausted queries, if available
        if config.QueryRemain > 0 {
            c.Logger.Debugf("ipstack: %v/%v, %v. geiop: %v/%v, %v. selected %v",
                c.Services["ipstack"].QueryRemain, c.Services["ipstack"].QueryLimit, c.Services["ipstack"].QueryLap,
                c.Services["geoip"].QueryRemain, c.Services["geoip"].QueryLimit, c.Services["geoip"].QueryLap,
                service)
            return service
        }
    }

    // if there are no services with unexhausted queries (i.e. QueryRemain == 0)
    // choose the one which made its last query earlier than others
    earliest := time.Now()
    for service, config := range c.Services {
        if config.QueryLap.Before(earliest) {
            result, earliest = service, config.QueryLap
        }
    }
    c.Logger.Debugf("ipstack: %v/%v, %v. geiop: %v/%v, %v. selected %v",
        c.Services["ipstack"].QueryRemain, c.Services["ipstack"].QueryLimit, c.Services["ipstack"].QueryLap,
        c.Services["geoip"].QueryRemain, c.Services["geoip"].QueryLimit, c.Services["geoip"].QueryLap,
        result)
    return result
}

// update - update the service's statistics about executed queries
func (c *client) update(service string) {

    c.Services[service].mux.Lock()
    c.Logger.Debugf("%v before update: %v/%v, %v",
        service, c.Services[service].QueryRemain, c.Services[service].QueryLimit, c.Services[service].QueryLap)

    c.Services[service].QueryLap = time.Now()

    if c.Services[service].QueryRemain > 0 {
        c.Services[service].QueryRemain -= 1
    }

    c.Logger.Debugf("%v after update: %v/%v, %v",
        service, c.Services[service].QueryRemain, c.Services[service].QueryLimit, c.Services[service].QueryLap)
    c.Services[service].mux.Unlock()
}

// answer - structure for unmarshalling  the servers JSON responses
type answer struct {
    Code    string   `json:"country_code"` // for response from `ipstack` service
    Country struct { // for response from `geoip` service
        Code string `json:"code"`
    } `json:"country"`
}

// query - gets a country code for the provided IP
func (c *client) query(IP string, serviceName string, result chan string, wg *sync.WaitGroup) {

    // TODO: only while debugging
    if c.Logger.Level == 5 && serviceName == "geoip" {
        time.Sleep(time.Second * 5)
    }
    runtime.Gosched() // TODO: only while debugging

    if wg != nil {
        defer wg.Done()
    }

    ur := url.URL{
        Host:   c.Services[serviceName].Host,
        Scheme: c.Services[serviceName].Scheme,
        Path:   c.Services[serviceName].Path + IP,
    }
    query := ur.Query()
    query.Set("access_key", c.Services[serviceName].Key)
    ur.RawQuery = query.Encode()
    c.Logger.Debugf("%v, query: %v", IP, ur.String())
    res, err := c.Client.Get(ur.String())
    defer res.Body.Close()
    if err != nil {
        c.Logger.Errorf("%v, failed to receive a response for the url %v: %v", IP, ur.String(), err)
    }

    var a answer
    decoder := json.NewDecoder(res.Body)
    err = decoder.Decode(&a)
    if err != nil {
        c.Logger.Errorf("%v, failed to unmarshal json for the url %v: %v", IP, ur.String(), err)
    }

    var code string
    switch serviceName {
    case "geoip":
        code = fmt.Sprintf("%v: %v (%v)", IP, a.Country.Code, serviceName)
    case "ipstack":
        code = fmt.Sprintf("%v: %v (%v)", IP, a.Code, serviceName)
    }
    c.Cache.Set(IP, code, cache.DefaultExpiration)
    c.Logger.Debugf("query result: %v", code)

    result <- code
    <-c.workersPull // release a worker
}

// Code - gets country codes for an array of IPs
func (c *client) Code(input []string, output chan string) {
    wg := &sync.WaitGroup{}
    for _, IP := range input {
        if code, found := c.Cache.Get(IP); found {
            c.Logger.Debugf("found in cache: %v", code)
            output <- code.(string)
            continue
        }
        wg.Add(1)
        serv := c.selectService()
        c.update(serv)

        c.workersPull <- struct{}{} // block if there are no free worker
        go c.query(IP, serv, output, wg)
    }
    wg.Wait()
    close(output)
}

Sample output

Full code + tests: https://github.com/OzRamsay/cquery

\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.