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)
}
Full code + tests: https://github.com/OzRamsay/cquery
