<commit_before>//===-- DWARFDebugRanges.cpp ------------------------------------*- C++ -*-===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#include "DWARFDebugRanges.h"
#include "DWARFUnit.h"
#include "SymbolFileDWARF.h"
#include "lldb/Utility/Stream.h"
#include <assert.h>

using namespace lldb_private;
using namespace std;

static dw_addr_t GetBaseAddressMarker(uint32_t addr_size) {
  switch(addr_size) {
    case 2:
      return 0xffff;
    case 4:
      return 0xffffffff;
    case 8:
      return 0xffffffffffffffff;
  }
  llvm_unreachable("GetBaseAddressMarker unsupported address size.");
}

DWARFDebugRanges::DWARFDebugRanges() : m_range_map() {}

void DWARFDebugRanges::Extract(SymbolFileDWARF *dwarf2Data) {
  DWARFRangeList range_list;
  lldb::offset_t offset = 0;
  dw_offset_t debug_ranges_offset = offset;
  while (Extract(dwarf2Data, &offset, range_list)) {
    range_list.Sort();
    m_range_map[debug_ranges_offset] = range_list;
    debug_ranges_offset = offset;
  }
}

bool DWARFDebugRanges::Extract(SymbolFileDWARF *dwarf2Data,
                               lldb::offset_t *offset_ptr,
                               DWARFRangeList &range_list) {
  range_list.Clear();

  lldb::offset_t range_offset = *offset_ptr;
  const DWARFDataExtractor &debug_ranges_data =
      dwarf2Data->get_debug_ranges_data();
  uint32_t addr_size = debug_ranges_data.GetAddressByteSize();
  dw_addr_t base_addr = 0;
  dw_addr_t base_addr_marker = GetBaseAddressMarker(addr_size);

  while (
      debug_ranges_data.ValidOffsetForDataOfSize(*offset_ptr, 2 * addr_size)) {
    dw_addr_t begin = debug_ranges_data.GetMaxU64(offset_ptr, addr_size);
    dw_addr_t end = debug_ranges_data.GetMaxU64(offset_ptr, addr_size);

    if (!begin && !end) {
      // End of range list
      break;
    }

    if (begin == base_addr_marker) {
      base_addr = end;
      continue;
    }

    // Filter out empty ranges
    if (begin < end)
      range_list.Append(DWARFRangeList::Entry(begin + base_addr, end - begin));
  }

  // Make sure we consumed at least something
  return range_offset != *offset_ptr;
}

void DWARFDebugRanges::Dump(Stream &s,
                            const DWARFDataExtractor &debug_ranges_data,
                            lldb::offset_t *offset_ptr,
                            dw_addr_t cu_base_addr) {
  uint32_t addr_size = s.GetAddressByteSize();

  dw_addr_t base_addr = cu_base_addr;
  while (
      debug_ranges_data.ValidOffsetForDataOfSize(*offset_ptr, 2 * addr_size)) {
    dw_addr_t begin = debug_ranges_data.GetMaxU64(offset_ptr, addr_size);
    dw_addr_t end = debug_ranges_data.GetMaxU64(offset_ptr, addr_size);
    // Extend 4 byte addresses that consists of 32 bits of 1's to be 64 bits of
    // ones
    if (begin == 0xFFFFFFFFull && addr_size == 4)
      begin = LLDB_INVALID_ADDRESS;

    s.Indent();
    if (begin == 0 && end == 0) {
      s.PutCString(" End");
      break;
    } else if (begin == LLDB_INVALID_ADDRESS) {
      // A base address selection entry
      base_addr = end;
      s.Address(base_addr, sizeof(dw_addr_t), " Base address = ");
    } else {
      // Convert from offset to an address
      dw_addr_t begin_addr = begin + base_addr;
      dw_addr_t end_addr = end + base_addr;

      s.AddressRange(begin_addr, end_addr, sizeof(dw_addr_t), NULL);
    }
  }
}

bool DWARFDebugRanges::FindRanges(const DWARFUnit *cu,
                                  dw_offset_t debug_ranges_offset,
                                  DWARFRangeList &range_list) const {
  dw_addr_t debug_ranges_address = cu->GetRangesBase() + debug_ranges_offset;
  range_map_const_iterator pos = m_range_map.find(debug_ranges_address);
  if (pos != m_range_map.end()) {
    range_list = pos->second;

    // All DW_AT_ranges are relative to the base address of the compile
    // unit. We add the compile unit base address to make sure all the
    // addresses are properly fixed up.
    range_list.Slide(cu->GetBaseAddress());
    return true;
  }
  return false;
}

bool DWARFDebugRngLists::ExtractRangeList(
    const DWARFDataExtractor &data, uint8_t addrSize,
    lldb::offset_t *offset_ptr, std::vector<RngListEntry> &rangeList) {
  rangeList.clear();

  bool error = false;
  while (!error) {
    switch (data.GetU8(offset_ptr)) {
    case DW_RLE_end_of_list:
      return true;

    case DW_RLE_start_length: {
      dw_addr_t begin = data.GetMaxU64(offset_ptr, addrSize);
      dw_addr_t len = data.GetULEB128(offset_ptr);
      rangeList.push_back({DW_RLE_start_length, begin, len});
      break;
    }

    case DW_RLE_start_end: {
      dw_addr_t begin = data.GetMaxU64(offset_ptr, addrSize);
      dw_addr_t end = data.GetMaxU64(offset_ptr, addrSize);
      rangeList.push_back({DW_RLE_start_end, begin, end});
      break;
    }

    case DW_RLE_base_address: {
      dw_addr_t base = data.GetMaxU64(offset_ptr, addrSize);
      rangeList.push_back({DW_RLE_base_address, base, 0});
      break;
    }

    case DW_RLE_offset_pair: {
      dw_addr_t begin = data.GetULEB128(offset_ptr);
      dw_addr_t end = data.GetULEB128(offset_ptr);
      rangeList.push_back({DW_RLE_offset_pair, begin, end});
      break;
    }

    default:
      // Next encodings are not yet supported:
      // DW_RLE_base_addressx, DW_RLE_startx_endx, DW_RLE_startx_length.
      lldbassert(0 && "unknown range list entry encoding");
      error = true;
    }
  }

  return false;
}

bool DWARFDebugRngLists::FindRanges(const DWARFUnit *cu,
                                    dw_offset_t debug_ranges_offset,
                                    DWARFRangeList &range_list) const {
  range_list.Clear();
  dw_addr_t debug_ranges_address = cu->GetRangesBase() + debug_ranges_offset;
  auto pos = m_range_map.find(debug_ranges_address);
  if (pos != m_range_map.end()) {
    dw_addr_t BaseAddr = cu->GetBaseAddress();
    for (const RngListEntry &E : pos->second) {
      switch (E.encoding) {
      case DW_RLE_start_length:
        range_list.Append(DWARFRangeList::Entry(E.value0, E.value1));
        break;
      case DW_RLE_base_address:
        BaseAddr = E.value0;
        break;
      case DW_RLE_start_end:
        range_list.Append(DWARFRangeList::Entry(E.value0, E.value1 - E.value0));
        break;
      case DW_RLE_offset_pair:
        range_list.Append(
            DWARFRangeList::Entry(BaseAddr + E.value0, E.value1 - E.value0));
        break;
      default:
        llvm_unreachable("unexpected encoding");
      }
    }
    return true;
  }
  return false;
}

void DWARFDebugRngLists::Extract(SymbolFileDWARF *dwarf2Data) {
  const DWARFDataExtractor &data = dwarf2Data->get_debug_rnglists_data();
  lldb::offset_t offset = 0;

  uint64_t length = data.GetU32(&offset);
  bool isDwarf64 = false;
  if (length == 0xffffffff) {
    length = data.GetU64(&offset);
    isDwarf64 = true;
  }
  lldb::offset_t end = offset + length;

  // Check version.
  if (data.GetU16(&offset) < 5)
    return;

  uint8_t addrSize = data.GetU8(&offset);

  // We do not support non-zero segment selector size.
  if (data.GetU8(&offset) != 0) {
    lldbassert(0 && "not implemented");
    return;
  }

  uint32_t offsetsAmount = data.GetU32(&offset);
  for (uint32_t i = 0; i < offsetsAmount; ++i)
    Offsets.push_back(data.GetPointer(&offset));

  lldb::offset_t listOffset = offset;
  std::vector<RngListEntry> rangeList;
  while (offset < end && ExtractRangeList(data, addrSize, &offset, rangeList)) {
    m_range_map[listOffset] = rangeList;
    listOffset = offset;
  }
}
<commit_msg>[LLDB] - Removed unused variable. NFC.<commit_after>//===-- DWARFDebugRanges.cpp ------------------------------------*- C++ -*-===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#include "DWARFDebugRanges.h"
#include "DWARFUnit.h"
#include "SymbolFileDWARF.h"
#include "lldb/Utility/Stream.h"
#include <assert.h>

using namespace lldb_private;
using namespace std;

static dw_addr_t GetBaseAddressMarker(uint32_t addr_size) {
  switch(addr_size) {
    case 2:
      return 0xffff;
    case 4:
      return 0xffffffff;
    case 8:
      return 0xffffffffffffffff;
  }
  llvm_unreachable("GetBaseAddressMarker unsupported address size.");
}

DWARFDebugRanges::DWARFDebugRanges() : m_range_map() {}

void DWARFDebugRanges::Extract(SymbolFileDWARF *dwarf2Data) {
  DWARFRangeList range_list;
  lldb::offset_t offset = 0;
  dw_offset_t debug_ranges_offset = offset;
  while (Extract(dwarf2Data, &offset, range_list)) {
    range_list.Sort();
    m_range_map[debug_ranges_offset] = range_list;
    debug_ranges_offset = offset;
  }
}

bool DWARFDebugRanges::Extract(SymbolFileDWARF *dwarf2Data,
                               lldb::offset_t *offset_ptr,
                               DWARFRangeList &range_list) {
  range_list.Clear();

  lldb::offset_t range_offset = *offset_ptr;
  const DWARFDataExtractor &debug_ranges_data =
      dwarf2Data->get_debug_ranges_data();
  uint32_t addr_size = debug_ranges_data.GetAddressByteSize();
  dw_addr_t base_addr = 0;
  dw_addr_t base_addr_marker = GetBaseAddressMarker(addr_size);

  while (
      debug_ranges_data.ValidOffsetForDataOfSize(*offset_ptr, 2 * addr_size)) {
    dw_addr_t begin = debug_ranges_data.GetMaxU64(offset_ptr, addr_size);
    dw_addr_t end = debug_ranges_data.GetMaxU64(offset_ptr, addr_size);

    if (!begin && !end) {
      // End of range list
      break;
    }

    if (begin == base_addr_marker) {
      base_addr = end;
      continue;
    }

    // Filter out empty ranges
    if (begin < end)
      range_list.Append(DWARFRangeList::Entry(begin + base_addr, end - begin));
  }

  // Make sure we consumed at least something
  return range_offset != *offset_ptr;
}

void DWARFDebugRanges::Dump(Stream &s,
                            const DWARFDataExtractor &debug_ranges_data,
                            lldb::offset_t *offset_ptr,
                            dw_addr_t cu_base_addr) {
  uint32_t addr_size = s.GetAddressByteSize();

  dw_addr_t base_addr = cu_base_addr;
  while (
      debug_ranges_data.ValidOffsetForDataOfSize(*offset_ptr, 2 * addr_size)) {
    dw_addr_t begin = debug_ranges_data.GetMaxU64(offset_ptr, addr_size);
    dw_addr_t end = debug_ranges_data.GetMaxU64(offset_ptr, addr_size);
    // Extend 4 byte addresses that consists of 32 bits of 1's to be 64 bits of
    // ones
    if (begin == 0xFFFFFFFFull && addr_size == 4)
      begin = LLDB_INVALID_ADDRESS;

    s.Indent();
    if (begin == 0 && end == 0) {
      s.PutCString(" End");
      break;
    } else if (begin == LLDB_INVALID_ADDRESS) {
      // A base address selection entry
      base_addr = end;
      s.Address(base_addr, sizeof(dw_addr_t), " Base address = ");
    } else {
      // Convert from offset to an address
      dw_addr_t begin_addr = begin + base_addr;
      dw_addr_t end_addr = end + base_addr;

      s.AddressRange(begin_addr, end_addr, sizeof(dw_addr_t), NULL);
    }
  }
}

bool DWARFDebugRanges::FindRanges(const DWARFUnit *cu,
                                  dw_offset_t debug_ranges_offset,
                                  DWARFRangeList &range_list) const {
  dw_addr_t debug_ranges_address = cu->GetRangesBase() + debug_ranges_offset;
  range_map_const_iterator pos = m_range_map.find(debug_ranges_address);
  if (pos != m_range_map.end()) {
    range_list = pos->second;

    // All DW_AT_ranges are relative to the base address of the compile
    // unit. We add the compile unit base address to make sure all the
    // addresses are properly fixed up.
    range_list.Slide(cu->GetBaseAddress());
    return true;
  }
  return false;
}

bool DWARFDebugRngLists::ExtractRangeList(
    const DWARFDataExtractor &data, uint8_t addrSize,
    lldb::offset_t *offset_ptr, std::vector<RngListEntry> &rangeList) {
  rangeList.clear();

  bool error = false;
  while (!error) {
    switch (data.GetU8(offset_ptr)) {
    case DW_RLE_end_of_list:
      return true;

    case DW_RLE_start_length: {
      dw_addr_t begin = data.GetMaxU64(offset_ptr, addrSize);
      dw_addr_t len = data.GetULEB128(offset_ptr);
      rangeList.push_back({DW_RLE_start_length, begin, len});
      break;
    }

    case DW_RLE_start_end: {
      dw_addr_t begin = data.GetMaxU64(offset_ptr, addrSize);
      dw_addr_t end = data.GetMaxU64(offset_ptr, addrSize);
      rangeList.push_back({DW_RLE_start_end, begin, end});
      break;
    }

    case DW_RLE_base_address: {
      dw_addr_t base = data.GetMaxU64(offset_ptr, addrSize);
      rangeList.push_back({DW_RLE_base_address, base, 0});
      break;
    }

    case DW_RLE_offset_pair: {
      dw_addr_t begin = data.GetULEB128(offset_ptr);
      dw_addr_t end = data.GetULEB128(offset_ptr);
      rangeList.push_back({DW_RLE_offset_pair, begin, end});
      break;
    }

    default:
      // Next encodings are not yet supported:
      // DW_RLE_base_addressx, DW_RLE_startx_endx, DW_RLE_startx_length.
      lldbassert(0 && "unknown range list entry encoding");
      error = true;
    }
  }

  return false;
}

bool DWARFDebugRngLists::FindRanges(const DWARFUnit *cu,
                                    dw_offset_t debug_ranges_offset,
                                    DWARFRangeList &range_list) const {
  range_list.Clear();
  dw_addr_t debug_ranges_address = cu->GetRangesBase() + debug_ranges_offset;
  auto pos = m_range_map.find(debug_ranges_address);
  if (pos != m_range_map.end()) {
    dw_addr_t BaseAddr = cu->GetBaseAddress();
    for (const RngListEntry &E : pos->second) {
      switch (E.encoding) {
      case DW_RLE_start_length:
        range_list.Append(DWARFRangeList::Entry(E.value0, E.value1));
        break;
      case DW_RLE_base_address:
        BaseAddr = E.value0;
        break;
      case DW_RLE_start_end:
        range_list.Append(DWARFRangeList::Entry(E.value0, E.value1 - E.value0));
        break;
      case DW_RLE_offset_pair:
        range_list.Append(
            DWARFRangeList::Entry(BaseAddr + E.value0, E.value1 - E.value0));
        break;
      default:
        llvm_unreachable("unexpected encoding");
      }
    }
    return true;
  }
  return false;
}

void DWARFDebugRngLists::Extract(SymbolFileDWARF *dwarf2Data) {
  const DWARFDataExtractor &data = dwarf2Data->get_debug_rnglists_data();
  lldb::offset_t offset = 0;

  uint64_t length = data.GetU32(&offset);
  if (length == 0xffffffff)
    length = data.GetU64(&offset);
  lldb::offset_t end = offset + length;

  // Check version.
  if (data.GetU16(&offset) < 5)
    return;

  uint8_t addrSize = data.GetU8(&offset);

  // We do not support non-zero segment selector size.
  if (data.GetU8(&offset) != 0) {
    lldbassert(0 && "not implemented");
    return;
  }

  uint32_t offsetsAmount = data.GetU32(&offset);
  for (uint32_t i = 0; i < offsetsAmount; ++i)
    Offsets.push_back(data.GetPointer(&offset));

  lldb::offset_t listOffset = offset;
  std::vector<RngListEntry> rangeList;
  while (offset < end && ExtractRangeList(data, addrSize, &offset, rangeList)) {
    m_range_map[listOffset] = rangeList;
    listOffset = offset;
  }
}
<|endoftext|>


<commit_before>/**
 * Copyright (c) 2016-present, Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "caffe2/operators/jsd_op.h"

namespace caffe2 {

namespace {

static constexpr float kLOG_THRESHOLD() {
  return 1e-20;
}

inline float logit(float p) {
  // it computes log(p / (1-p))
  // to avoid numeric issue, hard code p log(p) when p approaches 0
  float x = std::min(std::max(p, kLOG_THRESHOLD()), 1 - kLOG_THRESHOLD());
  return -log(1. / x - 1.);
}

inline float entropy(float p) {
  if (p < kLOG_THRESHOLD() || 1 - p < kLOG_THRESHOLD()) {
    return 0.;
  } else {
    float q = 1 - p;
    return -p * log(p) - q * log(q);
  }
}
} // namespace

template <>
bool BernoulliJSDOp<float, CPUContext>::RunOnDevice() {
  auto& X = Input(0); // predicted probabilities
  auto& T = Input(1); // target probabilities
  auto* L = Output(0); // JSD loss output
  int N = X.size();
  CAFFE_ENFORCE_EQ(T.size(), N);
  L->Resize(N);
  auto* x_data = X.data<float>();
  auto* t_data = T.data<float>();
  auto* l_data = L->mutable_data<float>();
  for (int i = 0; i < N; i++) {
    auto p_mdl = x_data[i];
    auto p_emp = t_data[i];
    auto p_avg = (p_mdl + p_emp) / 2.;
    auto jsd = entropy(p_avg) - (entropy(p_mdl) + entropy(p_emp)) / 2.;
    l_data[i] = jsd;
  }
  return true;
}

template <>
bool BernoulliJSDGradientOp<float, CPUContext>::RunOnDevice() {
  auto& go = Input(0);
  auto& X = Input(1);
  auto& T = Input(2);
  auto* gi = Output(0);
  int N = X.size();
  gi->Resize(N);
  auto* go_data = go.data<float>();
  auto* x_data = X.data<float>();
  auto* t_data = T.data<float>();
  auto* gi_data = gi->mutable_data<float>();
  for (int i = 0; i < N; i++) {
    auto p_mdl = x_data[i];
    auto p_emp = t_data[i];
    auto p_avg = (p_mdl + p_emp) / 2.;
    auto g_jsd = (logit(p_mdl) - logit(p_avg)) / 2.;
    gi_data[i] = go_data[i] * g_jsd;
  }
  return true;
}
REGISTER_CPU_OPERATOR(BernoulliJSD, BernoulliJSDOp<float, CPUContext>);
REGISTER_CPU_OPERATOR(
    BernoulliJSDGradient,
    BernoulliJSDGradientOp<float, CPUContext>);
OPERATOR_SCHEMA(BernoulliJSD)
    .NumInputs(2)
    .NumOutputs(1)
    .SetDoc(R"DOC(
Computes the Jensen-Shannon divergence (JSD) between two Bernoulli distributions
where each is parametrized by a single probability.
)DOC")
    .Input(0, "X", "array of probabilities for prediction")
    .Input(0, "T", "array of probabilities for target")
    .Output(0, "L", "array of JSD losses");
OPERATOR_SCHEMA(BernoulliJSDGradient).NumInputs(3).NumOutputs(1);

class GetBernoulliJSDGradient : public GradientMakerBase {
  using GradientMakerBase::GradientMakerBase;
  vector<OperatorDef> GetGradientDefs() override {
    return SingleGradientDef(
        "BernoulliJSDGradient",
        "",
        vector<string>{GO(0), I(0), I(1)},
        vector<string>{GI(0)});
  }
};
REGISTER_GRADIENT(BernoulliJSD, GetBernoulliJSDGradient);

} // namespace caffe2
<commit_msg>[GanH/Easy]Fix blob dim<commit_after>/**
 * Copyright (c) 2016-present, Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "caffe2/operators/jsd_op.h"

namespace caffe2 {

namespace {

static constexpr float kLOG_THRESHOLD() {
  return 1e-20;
}

inline float logit(float p) {
  // it computes log(p / (1-p))
  // to avoid numeric issue, hard code p log(p) when p approaches 0
  float x = std::min(std::max(p, kLOG_THRESHOLD()), 1 - kLOG_THRESHOLD());
  return -log(1. / x - 1.);
}

inline float entropy(float p) {
  if (p < kLOG_THRESHOLD() || 1 - p < kLOG_THRESHOLD()) {
    return 0.;
  } else {
    float q = 1 - p;
    return -p * log(p) - q * log(q);
  }
}
} // namespace

template <>
bool BernoulliJSDOp<float, CPUContext>::RunOnDevice() {
  auto& X = Input(0); // predicted probabilities
  auto& T = Input(1); // target probabilities
  auto* L = Output(0); // JSD loss output
  int N = X.size();
  CAFFE_ENFORCE_EQ(T.size(), N);
  L->ResizeLike(X);
  auto* x_data = X.data<float>();
  auto* t_data = T.data<float>();
  auto* l_data = L->mutable_data<float>();
  for (int i = 0; i < N; i++) {
    auto p_mdl = x_data[i];
    auto p_emp = t_data[i];
    auto p_avg = (p_mdl + p_emp) / 2.;
    auto jsd = entropy(p_avg) - (entropy(p_mdl) + entropy(p_emp)) / 2.;
    l_data[i] = jsd;
  }
  return true;
}

template <>
bool BernoulliJSDGradientOp<float, CPUContext>::RunOnDevice() {
  auto& go = Input(0);
  auto& X = Input(1);
  auto& T = Input(2);
  auto* gi = Output(0);
  int N = X.size();
  gi->ResizeLike(X);
  auto* go_data = go.data<float>();
  auto* x_data = X.data<float>();
  auto* t_data = T.data<float>();
  auto* gi_data = gi->mutable_data<float>();
  for (int i = 0; i < N; i++) {
    auto p_mdl = x_data[i];
    auto p_emp = t_data[i];
    auto p_avg = (p_mdl + p_emp) / 2.;
    auto g_jsd = (logit(p_mdl) - logit(p_avg)) / 2.;
    gi_data[i] = go_data[i] * g_jsd;
  }
  return true;
}
REGISTER_CPU_OPERATOR(BernoulliJSD, BernoulliJSDOp<float, CPUContext>);
REGISTER_CPU_OPERATOR(
    BernoulliJSDGradient,
    BernoulliJSDGradientOp<float, CPUContext>);
OPERATOR_SCHEMA(BernoulliJSD)
    .NumInputs(2)
    .NumOutputs(1)
    .SetDoc(R"DOC(
Computes the Jensen-Shannon divergence (JSD) between two Bernoulli distributions
where each is parametrized by a single probability.
)DOC")
    .Input(0, "X", "array of probabilities for prediction")
    .Input(0, "T", "array of probabilities for target")
    .Output(0, "L", "array of JSD losses");
OPERATOR_SCHEMA(BernoulliJSDGradient).NumInputs(3).NumOutputs(1);

class GetBernoulliJSDGradient : public GradientMakerBase {
  using GradientMakerBase::GradientMakerBase;
  vector<OperatorDef> GetGradientDefs() override {
    return SingleGradientDef(
        "BernoulliJSDGradient",
        "",
        vector<string>{GO(0), I(0), I(1)},
        vector<string>{GI(0)});
  }
};
REGISTER_GRADIENT(BernoulliJSD, GetBernoulliJSDGradient);

} // namespace caffe2
<|endoftext|>


<commit_before>/*
    Colour ASCII art rendering support for Crystal Space 3D library
    Copyright (C) 2005 by dr.W.C.A. Wijngaards
    Based on aalib canvas by Andrew Zabolotny <bit@eltech.ru>

    This library is free software; you can redistribute it and/or
    modify it under the terms of the GNU Library General Public
    License as published by the Free Software Foundation; either
    version 2 of the License, or (at your option) any later version.

    This library 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
    Library General Public License for more details.

    You should have received a copy of the GNU Library General Public
    License along with this library; if not, write to the Free
    Software Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
*/

#include "cssysdef.h"
#include <stdarg.h>
#include <stdlib.h>
#include <limits.h>
#include "csqint.h"

#include "csgeom/csrect.h"
#include "csgfx/rgbpixel.h"
#include "csutil/cfgacc.h"
#include "iutil/cfgfile.h"
#include "iutil/event.h"
#include "iutil/eventq.h"
#include "iutil/objreg.h"
#include "ivaria/reporter.h"

#include "cscaca.h"

//-------------------------------------------------------- csGraphics2DCaca ---

CS_IMPLEMENT_PLUGIN

SCF_IMPLEMENT_FACTORY (csGraphics2DCaca)

csGraphics2DCaca::csGraphics2DCaca (iBase *iParent) :
  scfImplementationType (this, iParent), cucul_canvas (0), dither (0),
  caca_display (0)
{
  EventOutlet = 0;
  Memory = 0;
}

csGraphics2DCaca::~csGraphics2DCaca (void)
{
  Close ();
}

bool csGraphics2DCaca::Initialize (iObjectRegistry *object_reg)
{
  if (!csGraphics2D::Initialize (object_reg))
    return false;

  csConfigAccess config;
  config.AddConfig(object_reg, "/config/cacacanvas.cfg");
  config.AddConfig(object_reg, "/config/video.cfg");

  // Load settings from config file and setup the aa_defparams structure
  HardwareCursor = config->GetBool ("Video.SystemMouseCursor", true);

  // set internal render format
  Depth = 32;
  pfmt.RedMask   = 0xff << 16;
  pfmt.GreenMask = 0xff << 8;
  pfmt.BlueMask  = 0xff;
  pfmt.AlphaMask = 0;
  pfmt.PalEntries = 0;
  pfmt.PixelBytes = 4;
  pfmt.complete ();
  _GetPixelAt = GetPixelAt32;

  // create event outlet
  csRef<iEventQueue> q = csQueryRegistry<iEventQueue> (object_reg);
  if (q != 0)
  {
    if (!EventOutlet.IsValid())
      EventOutlet = q->CreateEventOutlet (this);
  }
  return true;
}

bool csGraphics2DCaca::Open ()
{
  if (is_open) return true;

  csConfigAccess config;
  config.AddConfig(object_reg, "/config/cacacanvas.cfg");
  config.AddConfig(object_reg, "/config/video.cfg");

  if(config->KeyExists("Video.ASCII.Console.Size"))
    CS::Utility::setenv ("CACA_GEOMETRY",
      config->GetStr("Video.ASCII.Console.Size"), false);
  if(config->KeyExists("Video.ASCII.Console.Driver"))
    CS::Utility::setenv ("CACA_DRIVER",
      config->GetStr("Video.ASCII.Console.Driver", "x11"), false);
  if(config->KeyExists("Video.ASCII.Console.Font"))
    CS::Utility::setenv ("CACA_FONT",
      config->GetStr("Video.ASCII.Console.Font", "fixed"), false);

  fbWidth = config->GetInt("Video.Ascii.Offscreen.Width", 320);
  fbHeight = config->GetInt("Video.Ascii.Offscreen.Height", 240);

  cucul_canvas = cucul_create_canvas (0, 0);
  if (cucul_canvas == 0) return false;
  caca_display = caca_create_display (cucul_canvas);
  if (caca_display == 0)
    return false;
  cucul_clear_canvas (cucul_canvas);

  dither = cucul_create_dither (Depth, fbWidth, fbHeight,
    fbWidth*pfmt.PixelBytes, pfmt.RedMask, pfmt.GreenMask,
    pfmt.BlueMask, pfmt.AlphaMask);
  if (dither == 0)
    return false;
  if(config->KeyExists("Video.ASCII.Console.AntiAlias"))
    cucul_set_dither_antialias (dither,
      config->GetStr("Video.ASCII.Console.AntiAlias"));
  if(config->KeyExists("Video.ASCII.Console.Dither"))
    cucul_set_dither_mode (dither,
      config->GetStr("Video.ASCII.Console.Dither"));

  caca_set_display_title (caca_display, win_title);

  Memory = new unsigned char[ pfmt.PixelBytes*fbWidth*fbHeight ];

  return csGraphics2D::Open ();
}

void csGraphics2DCaca::Close ()
{
  if (dither)
  {
    cucul_free_dither (dither);
    dither = 0;
  }
  if(caca_display)
  {
    caca_free_display(caca_display);
    caca_display = 0;
  }
  if (cucul_canvas)
  {
    cucul_free_canvas (cucul_canvas);
    cucul_canvas = 0;
  }
  if(Memory) delete[] Memory;
  Memory = 0;

  if (!is_open) return;
  csGraphics2D::Close ();
}

int csGraphics2DCaca::MapKey(int raw)
{

  switch(raw)
  {
    case CACA_KEY_UNKNOWN: return 0;
    case CACA_KEY_BACKSPACE : return CSKEY_BACKSPACE;
    case CACA_KEY_TAB : return CSKEY_TAB;
    case CACA_KEY_RETURN : return CSKEY_ENTER;
    case CACA_KEY_PAUSE : return CSKEY_PAUSE;
    case CACA_KEY_ESCAPE : return CSKEY_ESC;
    case CACA_KEY_DELETE : return CSKEY_DEL;
    case CACA_KEY_UP : return CSKEY_UP;
    case CACA_KEY_DOWN : return CSKEY_DOWN;
    case CACA_KEY_LEFT : return CSKEY_LEFT;
    case CACA_KEY_RIGHT : return CSKEY_RIGHT;
    case CACA_KEY_INSERT : return CSKEY_INS;
    case CACA_KEY_HOME : return CSKEY_HOME;
    case CACA_KEY_END : return CSKEY_END;
    case CACA_KEY_PAGEUP : return CSKEY_PGUP;
    case CACA_KEY_PAGEDOWN : return CSKEY_PGDN;
    case CACA_KEY_F1 : return CSKEY_F1;
    case CACA_KEY_F2 : return CSKEY_F2;
    case CACA_KEY_F3 : return CSKEY_F3;
    case CACA_KEY_F4 : return CSKEY_F4;
    case CACA_KEY_F5 : return CSKEY_F5;
    case CACA_KEY_F6 : return CSKEY_F6;
    case CACA_KEY_F7 : return CSKEY_F7;
    case CACA_KEY_F8 : return CSKEY_F8;
    case CACA_KEY_F9 : return CSKEY_F9;
    case CACA_KEY_F10 : return CSKEY_F10;
    case CACA_KEY_F11 : return CSKEY_F11;
    case CACA_KEY_F12 : return CSKEY_F12;
    case CACA_KEY_F13 : return 0; // no CSKEY for this
    case CACA_KEY_F14 : return 0; // no CSKEY for this
    case CACA_KEY_F15 : return 0; // no CSKEY for this
    default:
      return raw;
  }
}

void csGraphics2DCaca::Print (csRect const* area)
{
  cucul_dither_bitmap (cucul_canvas, 0, 0,
    cucul_get_canvas_width (cucul_canvas),
    cucul_get_canvas_height (cucul_canvas), dither, Memory);
  caca_refresh_display (caca_display);

  /* Get all events from keyboard and mouse and put them into system queue */
  caca_event_t event;
  while (caca_get_event (caca_display, caca_event::CACA_EVENT_ANY, &event, 0))
  {
    switch(event.type)
    {
    case caca_event::CACA_EVENT_KEY_PRESS:
    case caca_event::CACA_EVENT_KEY_RELEASE:
      {
        utf32_char raw, cooked;
        if (event.data.key.utf32 != 0)
          raw = cooked = event.data.key.utf32;
        else
          raw = cooked = MapKey (event.data.key.ch);
	EventOutlet->Key (raw, cooked,
	  (event.type == caca_event::CACA_EVENT_KEY_PRESS));
      }
      break;
    case caca_event::CACA_EVENT_MOUSE_PRESS:
      EventOutlet->Mouse (event.data.mouse.button- 1, true,
        caca_get_mouse_x (caca_display), caca_get_mouse_y (caca_display));
      break;
    case caca_event::CACA_EVENT_MOUSE_RELEASE:
      EventOutlet->Mouse (event.data.mouse.button - 1, false,
        caca_get_mouse_x (caca_display), caca_get_mouse_y (caca_display));
      break;
    case caca_event::CACA_EVENT_MOUSE_MOTION:
      EventOutlet->Mouse (csmbNone, false,
        event.data.mouse.x, event.data.mouse.y);
      break;
    default:
      break;
    }
  }
}

bool csGraphics2DCaca::BeginDraw ()
{
  csGraphics2D::BeginDraw ();
  return true;
}

void csGraphics2DCaca::FinishDraw ()
{
  csGraphics2D::FinishDraw ();
}

void csGraphics2DCaca::SetTitle(const char* title)
{
  caca_set_display_title (caca_display, title);
}
<commit_msg>Missing include for setenv<commit_after>/*
    Colour ASCII art rendering support for Crystal Space 3D library
    Copyright (C) 2005 by dr.W.C.A. Wijngaards
    Based on aalib canvas by Andrew Zabolotny <bit@eltech.ru>

    This library is free software; you can redistribute it and/or
    modify it under the terms of the GNU Library General Public
    License as published by the Free Software Foundation; either
    version 2 of the License, or (at your option) any later version.

    This library 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
    Library General Public License for more details.

    You should have received a copy of the GNU Library General Public
    License along with this library; if not, write to the Free
    Software Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
*/

#include "cssysdef.h"
#include <stdarg.h>
#include <stdlib.h>
#include <limits.h>
#include "csqint.h"

#include "csgeom/csrect.h"
#include "csgfx/rgbpixel.h"
#include "csutil/cfgacc.h"
#include "csutil/setenv.h"
#include "iutil/cfgfile.h"
#include "iutil/event.h"
#include "iutil/eventq.h"
#include "iutil/objreg.h"
#include "ivaria/reporter.h"

#include "cscaca.h"

//-------------------------------------------------------- csGraphics2DCaca ---

CS_IMPLEMENT_PLUGIN

SCF_IMPLEMENT_FACTORY (csGraphics2DCaca)

csGraphics2DCaca::csGraphics2DCaca (iBase *iParent) :
  scfImplementationType (this, iParent), cucul_canvas (0), dither (0),
  caca_display (0)
{
  EventOutlet = 0;
  Memory = 0;
}

csGraphics2DCaca::~csGraphics2DCaca (void)
{
  Close ();
}

bool csGraphics2DCaca::Initialize (iObjectRegistry *object_reg)
{
  if (!csGraphics2D::Initialize (object_reg))
    return false;

  csConfigAccess config;
  config.AddConfig(object_reg, "/config/cacacanvas.cfg");
  config.AddConfig(object_reg, "/config/video.cfg");

  // Load settings from config file and setup the aa_defparams structure
  HardwareCursor = config->GetBool ("Video.SystemMouseCursor", true);

  // set internal render format
  Depth = 32;
  pfmt.RedMask   = 0xff << 16;
  pfmt.GreenMask = 0xff << 8;
  pfmt.BlueMask  = 0xff;
  pfmt.AlphaMask = 0;
  pfmt.PalEntries = 0;
  pfmt.PixelBytes = 4;
  pfmt.complete ();
  _GetPixelAt = GetPixelAt32;

  // create event outlet
  csRef<iEventQueue> q = csQueryRegistry<iEventQueue> (object_reg);
  if (q != 0)
  {
    if (!EventOutlet.IsValid())
      EventOutlet = q->CreateEventOutlet (this);
  }
  return true;
}

bool csGraphics2DCaca::Open ()
{
  if (is_open) return true;

  csConfigAccess config;
  config.AddConfig(object_reg, "/config/cacacanvas.cfg");
  config.AddConfig(object_reg, "/config/video.cfg");

  if(config->KeyExists("Video.ASCII.Console.Size"))
    CS::Utility::setenv ("CACA_GEOMETRY",
      config->GetStr("Video.ASCII.Console.Size"), false);
  if(config->KeyExists("Video.ASCII.Console.Driver"))
    CS::Utility::setenv ("CACA_DRIVER",
      config->GetStr("Video.ASCII.Console.Driver", "x11"), false);
  if(config->KeyExists("Video.ASCII.Console.Font"))
    CS::Utility::setenv ("CACA_FONT",
      config->GetStr("Video.ASCII.Console.Font", "fixed"), false);

  fbWidth = config->GetInt("Video.Ascii.Offscreen.Width", 320);
  fbHeight = config->GetInt("Video.Ascii.Offscreen.Height", 240);

  cucul_canvas = cucul_create_canvas (0, 0);
  if (cucul_canvas == 0) return false;
  caca_display = caca_create_display (cucul_canvas);
  if (caca_display == 0)
    return false;
  cucul_clear_canvas (cucul_canvas);

  dither = cucul_create_dither (Depth, fbWidth, fbHeight,
    fbWidth*pfmt.PixelBytes, pfmt.RedMask, pfmt.GreenMask,
    pfmt.BlueMask, pfmt.AlphaMask);
  if (dither == 0)
    return false;
  if(config->KeyExists("Video.ASCII.Console.AntiAlias"))
    cucul_set_dither_antialias (dither,
      config->GetStr("Video.ASCII.Console.AntiAlias"));
  if(config->KeyExists("Video.ASCII.Console.Dither"))
    cucul_set_dither_mode (dither,
      config->GetStr("Video.ASCII.Console.Dither"));

  caca_set_display_title (caca_display, win_title);

  Memory = new unsigned char[ pfmt.PixelBytes*fbWidth*fbHeight ];

  return csGraphics2D::Open ();
}

void csGraphics2DCaca::Close ()
{
  if (dither)
  {
    cucul_free_dither (dither);
    dither = 0;
  }
  if(caca_display)
  {
    caca_free_display(caca_display);
    caca_display = 0;
  }
  if (cucul_canvas)
  {
    cucul_free_canvas (cucul_canvas);
    cucul_canvas = 0;
  }
  if(Memory) delete[] Memory;
  Memory = 0;

  if (!is_open) return;
  csGraphics2D::Close ();
}

int csGraphics2DCaca::MapKey(int raw)
{

  switch(raw)
  {
    case CACA_KEY_UNKNOWN: return 0;
    case CACA_KEY_BACKSPACE : return CSKEY_BACKSPACE;
    case CACA_KEY_TAB : return CSKEY_TAB;
    case CACA_KEY_RETURN : return CSKEY_ENTER;
    case CACA_KEY_PAUSE : return CSKEY_PAUSE;
    case CACA_KEY_ESCAPE : return CSKEY_ESC;
    case CACA_KEY_DELETE : return CSKEY_DEL;
    case CACA_KEY_UP : return CSKEY_UP;
    case CACA_KEY_DOWN : return CSKEY_DOWN;
    case CACA_KEY_LEFT : return CSKEY_LEFT;
    case CACA_KEY_RIGHT : return CSKEY_RIGHT;
    case CACA_KEY_INSERT : return CSKEY_INS;
    case CACA_KEY_HOME : return CSKEY_HOME;
    case CACA_KEY_END : return CSKEY_END;
    case CACA_KEY_PAGEUP : return CSKEY_PGUP;
    case CACA_KEY_PAGEDOWN : return CSKEY_PGDN;
    case CACA_KEY_F1 : return CSKEY_F1;
    case CACA_KEY_F2 : return CSKEY_F2;
    case CACA_KEY_F3 : return CSKEY_F3;
    case CACA_KEY_F4 : return CSKEY_F4;
    case CACA_KEY_F5 : return CSKEY_F5;
    case CACA_KEY_F6 : return CSKEY_F6;
    case CACA_KEY_F7 : return CSKEY_F7;
    case CACA_KEY_F8 : return CSKEY_F8;
    case CACA_KEY_F9 : return CSKEY_F9;
    case CACA_KEY_F10 : return CSKEY_F10;
    case CACA_KEY_F11 : return CSKEY_F11;
    case CACA_KEY_F12 : return CSKEY_F12;
    case CACA_KEY_F13 : return 0; // no CSKEY for this
    case CACA_KEY_F14 : return 0; // no CSKEY for this
    case CACA_KEY_F15 : return 0; // no CSKEY for this
    default:
      return raw;
  }
}

void csGraphics2DCaca::Print (csRect const* area)
{
  cucul_dither_bitmap (cucul_canvas, 0, 0,
    cucul_get_canvas_width (cucul_canvas),
    cucul_get_canvas_height (cucul_canvas), dither, Memory);
  caca_refresh_display (caca_display);

  /* Get all events from keyboard and mouse and put them into system queue */
  caca_event_t event;
  while (caca_get_event (caca_display, caca_event::CACA_EVENT_ANY, &event, 0))
  {
    switch(event.type)
    {
    case caca_event::CACA_EVENT_KEY_PRESS:
    case caca_event::CACA_EVENT_KEY_RELEASE:
      {
        utf32_char raw, cooked;
        if (event.data.key.utf32 != 0)
          raw = cooked = event.data.key.utf32;
        else
          raw = cooked = MapKey (event.data.key.ch);
	EventOutlet->Key (raw, cooked,
	  (event.type == caca_event::CACA_EVENT_KEY_PRESS));
      }
      break;
    case caca_event::CACA_EVENT_MOUSE_PRESS:
      EventOutlet->Mouse (event.data.mouse.button- 1, true,
        caca_get_mouse_x (caca_display), caca_get_mouse_y (caca_display));
      break;
    case caca_event::CACA_EVENT_MOUSE_RELEASE:
      EventOutlet->Mouse (event.data.mouse.button - 1, false,
        caca_get_mouse_x (caca_display), caca_get_mouse_y (caca_display));
      break;
    case caca_event::CACA_EVENT_MOUSE_MOTION:
      EventOutlet->Mouse (csmbNone, false,
        event.data.mouse.x, event.data.mouse.y);
      break;
    default:
      break;
    }
  }
}

bool csGraphics2DCaca::BeginDraw ()
{
  csGraphics2D::BeginDraw ();
  return true;
}

void csGraphics2DCaca::FinishDraw ()
{
  csGraphics2D::FinishDraw ();
}

void csGraphics2DCaca::SetTitle(const char* title)
{
  caca_set_display_title (caca_display, title);
}
<|endoftext|>


<commit_before>#ifndef ALGORITHMIC_SKIP_LIST_HPP
#define ALGORITHMIC_SKIP_LIST_HPP
#include <algorithm>
#include <array>
#include <cassert>
#include <vector>

namespace algic
{
    using std::size_t;

    template <class T>
    struct node;

    struct node_base
    {
        node_base(size_t height);
        virtual ~node_base();

        size_t height() const;
        node_base* next(size_t level) const;
        void set(size_t level, node_base* nb);

        void incHeight();

        template <class T>
        node<T>* as()
        {
            return dynamic_cast<node<T>*>(this);
        }

    private:
        size_t  mHeight;
        node_base**  mLinks;
    };

    template <class T>
    struct node : public node_base
    {
        node(size_t height, T const& t);
        node(size_t height, T&& t);
        ~node();

        T& value();
        T const& value() const;

    private:
        T* asT();
        T const* asT() const;

        char mData[sizeof(T)];
    };

    /**********************************************************/
    /*                  implementations                       */

    /**********************************************************/
    /*                      node_base                         */
    node_base::node_base(size_t height)
        : mHeight(height)
        , mLinks(new node_base*[mHeight])
    {
        std::fill(mLinks, mLinks + mHeight, nullptr);
    }

    node_base::~node_base()
    {
        delete [] mLinks;
    }

    size_t node_base::height() const
    {
        return mHeight;
    }

    node_base* node_base::next(size_t level) const
    {
        assert(level < mHeight && "Index out of range");
        return mLinks[level];
    }

    void node_base::set(size_t level, node_base* nb)
    {
        assert(level < mHeight && "Index out of range");
        mLinks[level] = nb;
    }

    void node_base::incHeight()
    {
        ++mHeight;
        node_base** newLinks = new node_base*[mHeight];
        std::copy(mLinks, mLinks + mHeight, newLinks);
        delete mLinks;
        mLinks = newLinks;
        mLinks[mHeight - 1] = nullptr;
    }

    /**********************************************************/
    /*                         node                           */
    template <class T>
    node<T>::node(size_t height, T const& t)
        : node_base(height)
    {
        new (static_cast<void*>(mData))T(t);
    }

    template <class T>
    node<T>::node(size_t height, T&& t)
        : node_base(height)
    {
        new (static_cast<void*>(mData))T(std::move(t));
    }

    template <class T>
    node<T>::~node()
    {
        static_cast<T*>(static_cast<void*>(mData))->~T();
    }

    template <class T>
    T& node<T>::value()
    {
        return static_cast<T&>(*asT());
    }

    template <class T>
    T const& node<T>::value() const
    {
        return static_cast<T const&>(*asT());
    }

    template <class T>
    T* node<T>::asT()
    {
        return static_cast<T*>(static_cast<void*>(mData));
    }

    template <class T>
    T const* node<T>::asT() const
    {
        return static_cast<T const*>(static_cast<void*>(mData));
    }

    /**********************************************************/
    /*                      skip_list                         */
    template <class T, class RandomGen>
    skip_list<T, RandomGen>::skip_list(RandomGen& randGen, float prob)
        : mRand(randGen)
        , mHead(new node_base(1))
        , mProb(prob)
    {
        assert(prob >= 0.0f && "Probability can't be negative");
        assert(prob <= 1.0f && "Probability must be not greater than 1");
    }

    template <class T, class RandomGen>
    skip_list<T, RandomGen>::~skip_list()
    {
        delete mHead;
    }

    template <class T, class RandomGen>
    bool skip_list<T, RandomGen>::insert(T const& t)
    {
    }

    template <class T, class RandomGen>
    bool skip_list<T, RandomGen>::remove(T const& t)
    {
    }

    template <class T, class RandomGen>
    bool skip_list<T, RandomGen>::contains(T const& t)
    {
        size_t const H = mHead->height();
        for (size_t lvl = H - 1; lvl >= 0; --lvl)
        {
            if (node_base* place = findPlace(t, H, lvl, mHead->next(lvl)))
            {
                if (place->as<T>()->value() == t)
                    return true;
                else
                    return false;
            }
            if (lvl == 0)
                break;
        }
        return false;
    }

    template <class T, class RandomGen>
    node_base* skip_list<T, RandomGen>::findPlace(T const& t, size_t height, size_t level, node_base* start)
    {
        if (!start)
            return nullptr;

        auto const& val = start->as<T>()->value();
        if (val < t)
        {
            auto res = findPlace(t, height, level, start->next(level));
            if (!res)
            {
                if (level > 0)
                    return findPlace(t, height, level - 1, start);
                else
                    return start;
            }
        }
        if (val > t)
            return nullptr;
        return start;
    }

    template <class T, class RandomGen>
    bool skip_list<T, RandomGen>::coin() const
    {
        return mRand() < mProb;
    }

    template <class T, class RandomGen>
    size_t skip_list<T, RandomGen>::multiCoin() const
    {
        size_t count{ 0 };
        while (mRand() < mProb)
            ++count;
        return count;
    }
} // namespace algic

#endif
<commit_msg>implemented insert (does not work yet)<commit_after>#ifndef ALGORITHMIC_SKIP_LIST_HPP
#define ALGORITHMIC_SKIP_LIST_HPP
#include <algorithm>
#include <array>
#include <cassert>
#include <functional>
#include <list>
#include <vector>

namespace algic
{
    using std::size_t;

    template <class T>
    struct node;

    struct node_base
    {
        node_base(size_t height);
        virtual ~node_base();

        size_t height() const;
        node_base* next(size_t level) const;
        void set(size_t level, node_base* nb);

        void incHeight();

        template <class T>
        node<T>* as()
        {
            return dynamic_cast<node<T>*>(this);
        }

    private:
        size_t  mHeight;
        node_base**  mLinks;
    };

    template <class T>
    struct node : public node_base
    {
        node(size_t height, T const& t);
        node(size_t height, T&& t);
        ~node();

        T& value();
        T const& value() const;

    private:
        T* asT();
        T const* asT() const;

        char mData[sizeof(T)];
    };

    template <class T>
    bool less(node_base* node, T const& val)
    {
        if (auto tNode = node->as<T>())
            return std::less<T>()(tNode->value(), val);
        else
            return true;
    }

    template <class T>
    bool eq(node_base* node, T const& val)
    {
        if (auto tNode = node->as<T>())
            return std::equal_to<T>()(tNode->value(), val);
        else
            return false;
    }

    template <class T>
    bool gt(node_base* node, T const& val)
    {
        if (auto tNode = node->as<T>())
            return std::greater<T>()(tNode->value(), val);
        else
            return false;
    }

    /**********************************************************/
    /*                  implementations                       */

    /**********************************************************/
    /*                      node_base                         */
    node_base::node_base(size_t height)
        : mHeight(height)
        , mLinks(new node_base*[mHeight])
    {
        std::fill(mLinks, mLinks + mHeight, nullptr);
    }

    node_base::~node_base()
    {
        delete [] mLinks;
    }

    size_t node_base::height() const
    {
        return mHeight;
    }

    node_base* node_base::next(size_t level) const
    {
        assert(level < mHeight && "Index out of range");
        return mLinks[level];
    }

    void node_base::set(size_t level, node_base* nb)
    {
        assert(level < mHeight && "Index out of range");
        mLinks[level] = nb;
    }

    void node_base::incHeight()
    {
        ++mHeight;
        node_base** newLinks = new node_base*[mHeight];
        std::copy(mLinks, mLinks + mHeight, newLinks);
        delete mLinks;
        mLinks = newLinks;
        mLinks[mHeight - 1] = nullptr;
    }

    /**********************************************************/
    /*                         node                           */
    template <class T>
    node<T>::node(size_t height, T const& t)
        : node_base(height)
    {
        new (static_cast<void*>(mData))T(t);
    }

    template <class T>
    node<T>::node(size_t height, T&& t)
        : node_base(height)
    {
        new (static_cast<void*>(mData))T(std::move(t));
    }

    template <class T>
    node<T>::~node()
    {
        static_cast<T*>(static_cast<void*>(mData))->~T();
    }

    template <class T>
    T& node<T>::value()
    {
        return static_cast<T&>(*asT());
    }

    template <class T>
    T const& node<T>::value() const
    {
        return static_cast<T const&>(*asT());
    }

    template <class T>
    T* node<T>::asT()
    {
        return static_cast<T*>(static_cast<void*>(mData));
    }

    template <class T>
    T const* node<T>::asT() const
    {
        return static_cast<T const*>(static_cast<void*>(mData));
    }

    /**********************************************************/
    /*                      skip_list                         */
    template <class T, class RandomGen>
    skip_list<T, RandomGen>::skip_list(RandomGen& randGen, float prob)
        : mRand(randGen)
        , mHead(new node_base(1))
        , mProb(prob)
    {
        assert(prob >= 0.0f && "Probability can't be negative");
        assert(prob <= 1.0f && "Probability must be not greater than 1");
    }

    template <class T, class RandomGen>
    skip_list<T, RandomGen>::~skip_list()
    {
        delete mHead;
    }

    template <class T, class RandomGen>
    bool skip_list<T, RandomGen>::insert(T const& t)
    {
        size_t const H = mHead->height();
        std::vector<node_base*> visited(H, nullptr);

        node_base* curNode = mHead;
        int curLvl = H - 1;
        bool out = false;
        while (!out)
        {
            auto nextNode = curNode->next(curLvl);
            if (nextNode && less(nextNode, t))
            {
                curNode = nextNode;
                continue;
            }
            else if (curLvl >= 0)
            {
                visited[curLvl] = curNode;
                --curLvl;
                if (curLvl < 0)
                {
                    if (nextNode && eq(nextNode, t))
                        return false;
                    curLvl = 0;
                    out = true;
                    continue;
                }
            }
        }

        size_t const newLvl = multiCoin();
        node_base* newNode = new node<T>(newLvl, t);
        if (newLvl > H)
        {
            mHead->incHeight();
            mHead->set(H, newNode);
        }

        for (size_t i = 0; i < H - 1; ++i)
        {
            newNode->set(i, visited[i]);
            visited[i]->set(i, newNode);
        }
        return true;
    }

    template <class T, class RandomGen>
    bool skip_list<T, RandomGen>::remove(T const& t)
    {
    }

    template <class T, class RandomGen>
    bool skip_list<T, RandomGen>::contains(T const& t)
    {
        size_t const H = mHead->height();
        for (size_t lvl = H - 1; lvl >= 0; --lvl)
        {
            if (node_base* place = findPlace(t, H, lvl, mHead->next(lvl)))
            {
                if (place->as<T>()->value() == t)
                    return true;
                else
                    return false;
            }
            if (lvl == 0)
                break;
        }
        return false;
    }

    template <class T, class RandomGen>
    node_base* skip_list<T, RandomGen>::findPlace(T const& t, size_t height, size_t level, node_base* start)
    {
        if (!start)
            return nullptr;

        if (less(start, t))
        {
            auto res = findPlace(t, height, level, start->next(level));
            if (!res)
            {
                if (level > 0)
                    return findPlace(t, height, level - 1, start);
                else
                    return start;
            }
        }
        if (gt(start, t))
            return nullptr;
        return start;
    }

    template <class T, class RandomGen>
    bool skip_list<T, RandomGen>::coin() const
    {
        return mRand() < mProb;
    }

    template <class T, class RandomGen>
    size_t skip_list<T, RandomGen>::multiCoin() const
    {
        size_t const maxH{ mHead->height() + 1 };
        size_t count{ 0 };
        while (mRand() < mProb)
        {
            ++count;
            if (count == maxH)
                break;
        }
        return count;
    }
} // namespace algic

#endif
<|endoftext|>


<commit_before>#include "RpcImageServer.h"
#include "ImageServerCfg.h"
#include "ImageProcess.h"
#include "ImageProcessTask.h"
#include <image-manager/ImageServerStorage.h>
#include <image-manager/ImageServerRequest.h>
#include "errno.h"
#include <boost/bind.hpp>
#include <boost/filesystem.hpp>
#include <boost/thread.hpp>
#include <glog/logging.h>

#include <tfs_client_api.h>

using namespace tfs::client;
using namespace tfs::common;

namespace sf1r
{

RpcImageServer::RpcImageServer(const std::string& host, uint16_t port, uint32_t threadNum)
    : host_(host)
    , port_(port)
    , threadNum_(threadNum)
    , is_exporting_(false)
    , export_thread_(NULL)
{
}

RpcImageServer::~RpcImageServer()
{
    std::cout << "~RpcImageServer()" << std::endl;
    stop();
}

void RpcImageServer::start()
{
    instance.listen(host_, port_);
    instance.start(threadNum_);
}

void RpcImageServer::join()
{
    instance.join();
}

void RpcImageServer::run()
{
    start();
    join();
}

void RpcImageServer::stop()
{
    instance.end();
    instance.join();
    if(export_thread_)
    {
        export_thread_->interrupt();
        export_thread_->join();
        delete export_thread_;
    }
}

void RpcImageServer::dispatch(msgpack::rpc::request req)
{
    try
    {
        std::string method;
        req.method().convert(&method);

        if (method == ImageServerRequest::method_names[ImageServerRequest::METHOD_TEST])
        {
            msgpack::type::tuple<bool> params;
            req.params().convert(&params);

            req.result(true);
        }
        else if (method == ImageServerRequest::method_names[ImageServerRequest::METHOD_GET_IMAGE_COLOR])
        {
            msgpack::type::tuple<ImageColorData> params;
            req.params().convert(&params);
            ImageColorData& reqdata = params.get<0>();

            reqdata.success = ImageServerStorage::get()->GetImageColor(reqdata.img_file, reqdata.result_img_color);
            if(!reqdata.success)
            {
                // try to recalculate the color
                //std::string picFilePath = ImageServerCfg::get()->getImageFilePath() + "/" + reqdata.img_file;
                char* data = NULL;
                std::size_t datasize = 0;
                if(ImageServerStorage::get()->DownloadImageData(reqdata.img_file, data, datasize))
                {
                    int indexCollection[COLOR_RET_NUM];
                    bool ret = getImageColor(data, datasize, indexCollection, COLOR_RET_NUM);
                    if(ret){
                        std::ostringstream oss;
                        for(std::size_t i = 0; i < COLOR_RET_NUM - 1; ++i)
                        {
                            oss << indexCollection[i] << ",";
                        }
                        oss << indexCollection[COLOR_RET_NUM - 1];

                        reqdata.success = ImageServerStorage::get()->SetImageColor(reqdata.img_file, oss.str());
                        reqdata.result_img_color = oss.str();
                        //LOG(INFO) << "got request image color for fileName=" << reqdata.img_file << "result:" <<
                        //   oss.str() << std::endl;
                    }else{
                        LOG(WARNING) << "getImageColor failed imagePath=" << reqdata.img_file << std::endl;
                        reqdata.success = false;
                    }
                    delete[] data;
                }
                else
                {
                    LOG(INFO) << "getImage data failed while get image color imagePath=" << reqdata.img_file << std::endl;
                    reqdata.success = false;
                }
            }
            req.result(reqdata);
        }
        else if (method == ImageServerRequest::method_names[ImageServerRequest::METHOD_UPLOAD_IMAGE])
        {
            struct timeval tv_start, tv_end;
            gettimeofday(&tv_start, NULL);
            LOG(INFO) << "get upload request : " << tv_start.tv_sec << endl;
            msgpack::type::tuple<UploadImageData> params;
            req.params().convert(&params);
            UploadImageData& reqdata = params.get<0>();
            std::string ret_file;
            if(reqdata.param_type == 0)
            {
                reqdata.success = ImageServerStorage::get()->UploadImageFile(reqdata.img_file, ret_file);
            }
            else if(reqdata.param_type == 1)
            {
                reqdata.success = ImageServerStorage::get()->UploadImageData(reqdata.img_file.data(),
                    reqdata.img_file.size(), ret_file);
            }
            else
            {
                reqdata.success = false;
            }
            gettimeofday(&tv_end, NULL);
            int interval = tv_end.tv_sec - tv_start.tv_sec;
            if( interval > 1)
                LOG(WARNING) << "upload time larger than 1s : " << interval << endl;
            if(reqdata.success)
            {
                // write log file
                //
                ofstream ofs;
                try
                {
                    ofs.open(ImageServerCfg::get()->getUploadImageLog().c_str(), ofstream::app);
                    ofs << ret_file.c_str() << std::endl;
                    ofs.close();
                }
                catch(std::exception& e)
                {
                    LOG(ERROR) << "write upload log failed: " << ret_file << std::endl;
                    ofs.close();
                }
            }
            else
            {
                if(reqdata.param_type == 0)
                    LOG(WARNING) << "file upload failed: " << reqdata.img_file << std::endl;
                else
                    LOG(WARNING) << "image data upload failed: " << std::endl;
            }
            reqdata.img_file = ret_file;
            req.result(reqdata);
            gettimeofday(&tv_end, NULL);
            LOG(INFO) << "finish upload request : " << tv_end.tv_sec << endl;
            interval = tv_end.tv_sec - tv_start.tv_sec;
            if( interval > 1)
                LOG(WARNING) << "upload time total larger than 1s : " << interval << endl;
        }
        else if (method == ImageServerRequest::method_names[ImageServerRequest::METHOD_DELETE_IMAGE])
        {
            msgpack::type::tuple<DeleteImageData> params;
            req.params().convert(&params);
            DeleteImageData& reqdata = params.get<0>();
            reqdata.success = ImageServerStorage::get()->DeleteImage(reqdata.img_file);
            req.result(reqdata);
        }
        else if (method == ImageServerRequest::method_names[ImageServerRequest::METHOD_EXPORT_IMAGE])
        {
            msgpack::type::tuple<ExportImageData> params;
            req.params().convert(&params);
            ExportImageData& reqdata = params.get<0>();
            req.result(reqdata);
            if(is_exporting_)
            {
                LOG(WARNING) << "image is already exporting! Ignore any more request." << std::endl;
                return;
            }
            is_exporting_ = true;
            if(export_thread_)
            {
                delete export_thread_;
            }

            export_thread_ = new boost::thread(&RpcImageServer::exportimage, this,
                    reqdata.img_log_file, reqdata.out_dir);

        }
        else if(method == ImageServerRequest::method_names[ImageServerRequest::METHOD_COMPUTE_IMAGE_COLOR])
        {
            msgpack::type::tuple<ComputeImageColorData> params;
            req.params().convert(&params);
            ComputeImageColorData& reqdata = params.get<0>();

            Request *image_computereq = new Request;
            image_computereq->filePath = reqdata.filepath;
            image_computereq->filetype = reqdata.filetype;
            LOG(INFO) << "got image file color computing request:" << reqdata.filepath << endl;
            MSG msg(image_computereq);
            reqdata.success = true;
            if(!ImageProcessTask::instance()->put(msg)){
                LOG(ERROR) << "failed to enqueue fileName:" << reqdata.filepath << endl;
                reqdata.success = false;
            }
            req.result(reqdata);
        }
        else
        {
            req.error(msgpack::rpc::NO_METHOD_ERROR);
        }
    }
    catch (const msgpack::type_error& e)
    {
        req.error(msgpack::rpc::ARGUMENT_ERROR);
    }
    catch (const std::exception& e)
    {
        req.error(std::string(e.what()));
    }
}

void RpcImageServer::exportimage(const std::string& log_file, const std::string& outdir)
{
    size_t total_export = 0;
    size_t failed_cnt = 0;
    ifstream ifs;
    ofstream export_log;
    struct timeval tv_start, tv_end;
    try
    {
        ifs.open(log_file.c_str(), ifstream::in);
        export_log.open((outdir + "/" + std::string("export.log")).c_str(), ofstream::app);

        gettimeofday(&tv_start, NULL);
        export_log << " === export start: " << tv_start.tv_sec << " === " << std::endl;
        while(ifs.good())
        {
            ++total_export;
            char file_name[TFS_FILE_LEN];
            ifs.getline(file_name, TFS_FILE_LEN);
            char* buffer = NULL;
            std::size_t buf_size = 0;
            bool ret = ImageServerStorage::get()->DownloadImageData(file_name, buffer, buf_size);
            if(ret)
            {
                ofstream ofs;
                try
                {
                    ofs.open((outdir + "/" + file_name).c_str(), ofstream::trunc);
                    ofs.write(buffer, buf_size);
                }
                catch(std::exception& e)
                {
                    ++failed_cnt;
                    export_log << "export failed (save data to out file error) : " << e.what() << ",file: " << file_name << std::endl;
                }
                if(ofs.is_open())
                    ofs.close();
                delete[] buffer;
            }
            else
            {
                ++failed_cnt;
                export_log << "export failed (download error): " << file_name << std::endl;
            }
            boost::this_thread::interruption_point();
        }
    }
    catch(boost::thread_interrupted& )
    {
        LOG(INFO) << "exporting aborted. " << std::endl;
    }
    catch(std::exception& e)
    {
        LOG(WARNING) << "exception while exporting: " << e.what() << std::endl;
    }
    if(ifs.is_open())
    {
        ifs.close();
    }
    if(export_log.is_open())
    {
        export_log << "export finished, total export : " << total_export << ", failed: " << failed_cnt << std::endl;
        gettimeofday(&tv_end, NULL);
        export_log << " === export end: " << tv_end.tv_sec << " === " << std::endl;
        export_log.close();
    }
    is_exporting_ = false;
}

}
<commit_msg>tiny log for image server<commit_after>#include "RpcImageServer.h"
#include "ImageServerCfg.h"
#include "ImageProcess.h"
#include "ImageProcessTask.h"
#include <image-manager/ImageServerStorage.h>
#include <image-manager/ImageServerRequest.h>
#include "errno.h"
#include <boost/bind.hpp>
#include <boost/filesystem.hpp>
#include <boost/thread.hpp>
#include <glog/logging.h>

#include <tfs_client_api.h>

using namespace tfs::client;
using namespace tfs::common;

namespace sf1r
{

RpcImageServer::RpcImageServer(const std::string& host, uint16_t port, uint32_t threadNum)
    : host_(host)
    , port_(port)
    , threadNum_(threadNum)
    , is_exporting_(false)
    , export_thread_(NULL)
{
}

RpcImageServer::~RpcImageServer()
{
    std::cout << "~RpcImageServer()" << std::endl;
    stop();
}

void RpcImageServer::start()
{
    instance.listen(host_, port_);
    instance.start(threadNum_);
}

void RpcImageServer::join()
{
    instance.join();
}

void RpcImageServer::run()
{
    start();
    join();
}

void RpcImageServer::stop()
{
    instance.end();
    instance.join();
    if(export_thread_)
    {
        export_thread_->interrupt();
        export_thread_->join();
        delete export_thread_;
    }
}

void RpcImageServer::dispatch(msgpack::rpc::request req)
{
    try
    {
        std::string method;
        req.method().convert(&method);

        if (method == ImageServerRequest::method_names[ImageServerRequest::METHOD_TEST])
        {
            msgpack::type::tuple<bool> params;
            req.params().convert(&params);

            req.result(true);
        }
        else if (method == ImageServerRequest::method_names[ImageServerRequest::METHOD_GET_IMAGE_COLOR])
        {
            msgpack::type::tuple<ImageColorData> params;
            req.params().convert(&params);
            ImageColorData& reqdata = params.get<0>();

            reqdata.success = ImageServerStorage::get()->GetImageColor(reqdata.img_file, reqdata.result_img_color);
            if(!reqdata.success)
            {
                // try to recalculate the color
                //std::string picFilePath = ImageServerCfg::get()->getImageFilePath() + "/" + reqdata.img_file;
                char* data = NULL;
                std::size_t datasize = 0;
                if(ImageServerStorage::get()->DownloadImageData(reqdata.img_file, data, datasize))
                {
                    int indexCollection[COLOR_RET_NUM];
                    bool ret = getImageColor(data, datasize, indexCollection, COLOR_RET_NUM);
                    if(ret){
                        std::ostringstream oss;
                        for(std::size_t i = 0; i < COLOR_RET_NUM - 1; ++i)
                        {
                            oss << indexCollection[i] << ",";
                        }
                        oss << indexCollection[COLOR_RET_NUM - 1];

                        reqdata.success = ImageServerStorage::get()->SetImageColor(reqdata.img_file, oss.str());
                        reqdata.result_img_color = oss.str();
                        //LOG(INFO) << "got request image color for fileName=" << reqdata.img_file << "result:" <<
                        //   oss.str() << std::endl;
                    }else{
                        LOG(WARNING) << "getImageColor failed imagePath=" << reqdata.img_file << std::endl;
                        reqdata.success = false;
                    }
                    delete[] data;
                }
                else
                {
                    LOG(INFO) << "getImage data failed while get image color imagePath=" << reqdata.img_file << std::endl;
                    reqdata.success = false;
                }
            }
            req.result(reqdata);
        }
        else if (method == ImageServerRequest::method_names[ImageServerRequest::METHOD_UPLOAD_IMAGE])
        {
            struct timeval tv_start, tv_end;
            gettimeofday(&tv_start, NULL);
            //LOG(INFO) << "get upload request : " << tv_start.tv_sec << endl;
            msgpack::type::tuple<UploadImageData> params;
            req.params().convert(&params);
            UploadImageData& reqdata = params.get<0>();
            std::string ret_file;
            if(reqdata.param_type == 0)
            {
                reqdata.success = ImageServerStorage::get()->UploadImageFile(reqdata.img_file, ret_file);
            }
            else if(reqdata.param_type == 1)
            {
                reqdata.success = ImageServerStorage::get()->UploadImageData(reqdata.img_file.data(),
                    reqdata.img_file.size(), ret_file);
            }
            else
            {
                reqdata.success = false;
            }
            gettimeofday(&tv_end, NULL);
            int interval = tv_end.tv_sec - tv_start.tv_sec;
            if( interval > 1)
                LOG(WARNING) << "upload time larger than 1s : " << interval << endl;
            if(reqdata.success)
            {
                // write log file
                //
                ofstream ofs;
                try
                {
                    ofs.open(ImageServerCfg::get()->getUploadImageLog().c_str(), ofstream::app);
                    ofs << ret_file.c_str() << ", time: " << tv_end.tv_sec << std::endl;
                    ofs.close();
                }
                catch(std::exception& e)
                {
                    LOG(ERROR) << "write upload log failed: " << ret_file << std::endl;
                    ofs.close();
                }
            }
            else
            {
                if(reqdata.param_type == 0)
                    LOG(WARNING) << "file upload failed: " << reqdata.img_file << std::endl;
                else
                    LOG(WARNING) << "image data upload failed: " << std::endl;
            }
            reqdata.img_file = ret_file;
            req.result(reqdata);
            gettimeofday(&tv_end, NULL);
            //LOG(INFO) << "finish upload request : " << tv_end.tv_sec << endl;
            interval = tv_end.tv_sec - tv_start.tv_sec;
            if( interval > 1)
                LOG(WARNING) << "upload time total larger than 1s : " << interval << endl;
        }
        else if (method == ImageServerRequest::method_names[ImageServerRequest::METHOD_DELETE_IMAGE])
        {
            msgpack::type::tuple<DeleteImageData> params;
            req.params().convert(&params);
            DeleteImageData& reqdata = params.get<0>();
            reqdata.success = ImageServerStorage::get()->DeleteImage(reqdata.img_file);
            LOG(INFO) << "delete request for : " << reqdata.img_file << std::endl;
            req.result(reqdata);
        }
        else if (method == ImageServerRequest::method_names[ImageServerRequest::METHOD_EXPORT_IMAGE])
        {
            msgpack::type::tuple<ExportImageData> params;
            req.params().convert(&params);
            ExportImageData& reqdata = params.get<0>();
            req.result(reqdata);
            if(is_exporting_)
            {
                LOG(WARNING) << "image is already exporting! Ignore any more request." << std::endl;
                return;
            }
            is_exporting_ = true;
            if(export_thread_)
            {
                delete export_thread_;
            }

            export_thread_ = new boost::thread(&RpcImageServer::exportimage, this,
                    reqdata.img_log_file, reqdata.out_dir);

        }
        else if(method == ImageServerRequest::method_names[ImageServerRequest::METHOD_COMPUTE_IMAGE_COLOR])
        {
            msgpack::type::tuple<ComputeImageColorData> params;
            req.params().convert(&params);
            ComputeImageColorData& reqdata = params.get<0>();

            Request *image_computereq = new Request;
            image_computereq->filePath = reqdata.filepath;
            image_computereq->filetype = reqdata.filetype;
            LOG(INFO) << "got image file color computing request:" << reqdata.filepath << endl;
            MSG msg(image_computereq);
            reqdata.success = true;
            if(!ImageProcessTask::instance()->put(msg)){
                LOG(ERROR) << "failed to enqueue fileName:" << reqdata.filepath << endl;
                reqdata.success = false;
            }
            req.result(reqdata);
        }
        else
        {
            req.error(msgpack::rpc::NO_METHOD_ERROR);
        }
    }
    catch (const msgpack::type_error& e)
    {
        req.error(msgpack::rpc::ARGUMENT_ERROR);
    }
    catch (const std::exception& e)
    {
        req.error(std::string(e.what()));
    }
}

void RpcImageServer::exportimage(const std::string& log_file, const std::string& outdir)
{
    size_t total_export = 0;
    size_t failed_cnt = 0;
    ifstream ifs;
    ofstream export_log;
    struct timeval tv_start, tv_end;
    try
    {
        ifs.open(log_file.c_str(), ifstream::in);
        export_log.open((outdir + "/" + std::string("export.log")).c_str(), ofstream::app);

        gettimeofday(&tv_start, NULL);
        export_log << " === export start: " << tv_start.tv_sec << " === " << std::endl;
        while(ifs.good())
        {
            ++total_export;
            char file_name[TFS_FILE_LEN];
            ifs.getline(file_name, TFS_FILE_LEN);
            char* buffer = NULL;
            std::size_t buf_size = 0;
            bool ret = ImageServerStorage::get()->DownloadImageData(file_name, buffer, buf_size);
            if(ret)
            {
                ofstream ofs;
                try
                {
                    ofs.open((outdir + "/" + file_name).c_str(), ofstream::trunc);
                    ofs.write(buffer, buf_size);
                }
                catch(std::exception& e)
                {
                    ++failed_cnt;
                    export_log << "export failed (save data to out file error) : " << e.what() << ",file: " << file_name << std::endl;
                }
                if(ofs.is_open())
                    ofs.close();
                delete[] buffer;
            }
            else
            {
                ++failed_cnt;
                export_log << "export failed (download error): " << file_name << std::endl;
            }
            boost::this_thread::interruption_point();
        }
    }
    catch(boost::thread_interrupted& )
    {
        LOG(INFO) << "exporting aborted. " << std::endl;
    }
    catch(std::exception& e)
    {
        LOG(WARNING) << "exception while exporting: " << e.what() << std::endl;
    }
    if(ifs.is_open())
    {
        ifs.close();
    }
    if(export_log.is_open())
    {
        export_log << "export finished, total export : " << total_export << ", failed: " << failed_cnt << std::endl;
        gettimeofday(&tv_end, NULL);
        export_log << " === export end: " << tv_end.tv_sec << " === " << std::endl;
        export_log.close();
    }
    is_exporting_ = false;
}

}
<|endoftext|>


<commit_before>// Copyright 2012 Cloudera Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "util/webserver.h"

#include <boost/algorithm/string.hpp>
#include <boost/bind.hpp>
#include <boost/filesystem.hpp>
#include <boost/foreach.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/mem_fn.hpp>
#include <boost/thread/locks.hpp>
#include <gutil/strings/substitute.h>
#include <map>
#include <fstream>
#include <stdio.h>
#include <signal.h>
#include <string>
#include <mustache/mustache.h>
#include <rapidjson/stringbuffer.h>
#include <rapidjson/prettywriter.h>

#include "common/logging.h"
#include "util/cpu-info.h"
#include "util/disk-info.h"
#include "util/mem-info.h"
#include "util/os-info.h"
#include "util/url-coding.h"
#include "util/debug-util.h"
#include "util/stopwatch.h"
#include "rpc/thrift-util.h"

using namespace std;
using namespace boost;
using namespace boost::filesystem;
using namespace google;
using namespace strings;
using namespace rapidjson;
using namespace mustache;

const char* GetDefaultDocumentRoot();

DEFINE_int32(webserver_port, 25000, "Port to start debug webserver on");
DEFINE_string(webserver_interface, "",
    "Interface to start debug webserver on. If blank, webserver binds to 0.0.0.0");
DEFINE_string(webserver_doc_root, GetDefaultDocumentRoot(),
    "Files under <webserver_doc_root>/www are accessible via the debug webserver. "
    "Defaults to $IMPALA_HOME, or if $IMPALA_HOME is not set, disables the document "
    "root");
DEFINE_bool(enable_webserver_doc_root, true,
    "If true, webserver may serve static files from the webserver_doc_root");

DEFINE_string(webserver_certificate_file, "",
    "The location of the debug webserver's SSL certificate file, in .pem format. If "
    "empty, webserver SSL support is not enabled");
DEFINE_string(webserver_authentication_domain, "",
    "Domain used for debug webserver authentication");
DEFINE_string(webserver_password_file, "",
    "(Optional) Location of .htpasswd file containing user names and hashed passwords for"
    " debug webserver authentication");

static const char* DOC_FOLDER = "/www/";
static const int DOC_FOLDER_LEN = strlen(DOC_FOLDER);

// Easy-to-read constants for Squeasel return codes
static const uint32_t PROCESSING_COMPLETE = 1;
static const uint32_t NOT_PROCESSED = 0;

// Standard keys in the json document sent to templates for rendering. Must be kept in
// sync with the templates themselves.
static const char* COMMON_JSON_KEY = "__common__";

// Returns $IMPALA_HOME if set, otherwise /tmp/impala_www
const char* GetDefaultDocumentRoot() {
  stringstream ss;
  char* impala_home = getenv("IMPALA_HOME");
  if (impala_home == NULL) {
    return ""; // Empty document root means don't serve static files
  } else {
    ss << impala_home;
  }

  // Deliberate memory leak, but this should be called exactly once.
  string* str = new string(ss.str());
  return str->c_str();
}

namespace impala {

const char* Webserver::ENABLE_RAW_JSON_KEY = "__raw__";

Webserver::Webserver() : context_(NULL) {
  http_address_ = MakeNetworkAddress(
      FLAGS_webserver_interface.empty() ? "0.0.0.0" : FLAGS_webserver_interface,
      FLAGS_webserver_port);
}

Webserver::Webserver(const int port) : context_(NULL) {
  http_address_ = MakeNetworkAddress("0.0.0.0", port);
}

Webserver::~Webserver() {
  Stop();
}

void Webserver::RootHandler(const Webserver::ArgumentMap& args, Document* document) {
  document->AddMember("version", GetVersionString().c_str(),
      document->GetAllocator());
  document->AddMember("cpu_info", CpuInfo::DebugString().c_str(),
      document->GetAllocator());
  document->AddMember("mem_info", MemInfo::DebugString().c_str(),
      document->GetAllocator());
  document->AddMember("disk_info", DiskInfo::DebugString().c_str(),
      document->GetAllocator());
  document->AddMember("os_info", OsInfo::DebugString().c_str(),
      document->GetAllocator());
  document->AddMember("pid", getpid(), document->GetAllocator());
}

void Webserver::BuildArgumentMap(const string& args, ArgumentMap* output) {
  vector<string> arg_pairs;
  split(arg_pairs, args, is_any_of("&"));

  BOOST_FOREACH(const string& arg_pair, arg_pairs) {
    vector<string> key_value;
    split(key_value, arg_pair, is_any_of("="));
    if (key_value.empty()) continue;

    string key;
    if (!UrlDecode(key_value[0], &key)) continue;
    string value;
    if (!UrlDecode((key_value.size() >= 2 ? key_value[1] : ""), &value)) continue;
    to_lower(key);
    (*output)[key] = value;
  }
}

bool Webserver::IsSecure() const {
  return !FLAGS_webserver_certificate_file.empty();
}

Status Webserver::Start() {
  LOG(INFO) << "Starting webserver on " << http_address_;

  stringstream listening_spec;
  listening_spec << http_address_;

  if (IsSecure()) {
    LOG(INFO) << "Webserver: Enabling HTTPS support";
    // Squeasel makes sockets with 's' suffixes accept SSL traffic only
    listening_spec << "s";
  }
  string listening_str = listening_spec.str();
  vector<const char*> options;

  if (!FLAGS_webserver_doc_root.empty() && FLAGS_enable_webserver_doc_root) {
    LOG(INFO) << "Document root: " << FLAGS_webserver_doc_root;
    options.push_back("document_root");
    options.push_back(FLAGS_webserver_doc_root.c_str());
  } else {
    LOG(INFO)<< "Document root disabled";
  }

  if (IsSecure()) {
    options.push_back("ssl_certificate");
    options.push_back(FLAGS_webserver_certificate_file.c_str());
  }

  if (!FLAGS_webserver_authentication_domain.empty()) {
    options.push_back("authentication_domain");
    options.push_back(FLAGS_webserver_authentication_domain.c_str());
  }

  if (!FLAGS_webserver_password_file.empty()) {
    // Squeasel doesn't log anything if it can't stat the password file (but will if it
    // can't open it, which it tries to do during a request)
    if (!exists(FLAGS_webserver_password_file)) {
      stringstream ss;
      ss << "Webserver: Password file does not exist: " << FLAGS_webserver_password_file;
      return Status(ss.str());
    }
    LOG(INFO) << "Webserver: Password file is " << FLAGS_webserver_password_file;
    options.push_back("global_auth_file");
    options.push_back(FLAGS_webserver_password_file.c_str());
  }

  options.push_back("listening_ports");
  options.push_back(listening_str.c_str());

  // Options must be a NULL-terminated list
  options.push_back(NULL);

  // squeasel ignores SIGCHLD and we need it to run kinit. This means that since
  // squeasel does not reap its own children CGI programs must be avoided.
  // Save the signal handler so we can restore it after squeasel sets it to be ignored.
  sighandler_t sig_chld = signal(SIGCHLD, SIG_DFL);

  sq_callbacks callbacks;
  memset(&callbacks, 0, sizeof(callbacks));
  callbacks.begin_request = &Webserver::BeginRequestCallbackStatic;
  callbacks.log_message = &Webserver::LogMessageCallbackStatic;

  // To work around not being able to pass member functions as C callbacks, we store a
  // pointer to this server in the per-server state, and register a static method as the
  // default callback. That method unpacks the pointer to this and calls the real
  // callback.
  context_ = sq_start(&callbacks, reinterpret_cast<void*>(this), &options[0]);

  // Restore the child signal handler so wait() works properly.
  signal(SIGCHLD, sig_chld);

  if (context_ == NULL) {
    stringstream error_msg;
    error_msg << "Webserver: Could not start on address " << http_address_;
    return Status(error_msg.str());
  }

  UrlCallback default_callback =
      bind<void>(mem_fn(&Webserver::RootHandler), this, _1, _2);

  RegisterUrlCallback("/", "root.tmpl", default_callback);

  LOG(INFO) << "Webserver started";
  return Status::OK;
}

void Webserver::Stop() {
  if (context_ != NULL) {
    sq_stop(context_);
    context_ = NULL;
  }
}

void Webserver::GetCommonJson(Document* document) {
  DCHECK(document != NULL);
  Value obj(kObjectType);
  obj.AddMember("process-name", google::ProgramInvocationShortName(),
      document->GetAllocator());

  Value lst(kArrayType);
  BOOST_FOREACH(const UrlHandlerMap::value_type& handler, url_handlers_) {
    if (handler.second.is_on_nav_bar()) {
      Value obj(kObjectType);
      obj.AddMember("link", handler.first.c_str(), document->GetAllocator());
      obj.AddMember("title", handler.first.c_str(), document->GetAllocator());
      lst.PushBack(obj, document->GetAllocator());
    }
  }

  obj.AddMember("navbar", lst, document->GetAllocator());
  document->AddMember(COMMON_JSON_KEY, obj, document->GetAllocator());
}

int Webserver::LogMessageCallbackStatic(const struct sq_connection* connection,
    const char* message) {
  if (message != NULL) {
    LOG(INFO) << "Webserver: " << message;
  }
  return PROCESSING_COMPLETE;
}

int Webserver::BeginRequestCallbackStatic(struct sq_connection* connection) {
  struct sq_request_info* request_info = sq_get_request_info(connection);
  Webserver* instance = reinterpret_cast<Webserver*>(request_info->user_data);
  return instance->BeginRequestCallback(connection, request_info);
}

int Webserver::BeginRequestCallback(struct sq_connection* connection,
    struct sq_request_info* request_info) {
  if (!FLAGS_webserver_doc_root.empty() && FLAGS_enable_webserver_doc_root) {
    if (strncmp(DOC_FOLDER, request_info->uri, DOC_FOLDER_LEN) == 0) {
      VLOG(2) << "HTTP File access: " << request_info->uri;
      // Let Squeasel deal with this request; returning NULL will fall through
      // to the default handler which will serve files.
      return NOT_PROCESSED;
    }
  }
  shared_lock<shared_mutex> lock(url_handlers_lock_);
  UrlHandlerMap::const_iterator it = url_handlers_.find(request_info->uri);
  if (it == url_handlers_.end()) {
    sq_printf(connection, "HTTP/1.1 404 Not Found\r\n"
        "Content-Type: text/plain\r\n\r\n");
    sq_printf(connection, "No handler for URI %s\r\n\r\n", request_info->uri);
    return PROCESSING_COMPLETE;
  }

  map<string, string> arguments;
  if (request_info->query_string != NULL) {
    BuildArgumentMap(request_info->query_string, &arguments);
  }
  MonotonicStopWatch sw;
  sw.Start();

  Document document;
  document.SetObject();
  GetCommonJson(&document);

  bool send_html_headers = true;
  // The output of this page is accumulated into this stringstream.
  stringstream output;
  bool raw_json = (arguments.find("json") != arguments.end());
  if (raw_json) {
    // Callbacks may optionally be rendered as a text-only, pretty-printed Json document
    // (mostly for debugging or integration with third-party tools).
    StringBuffer strbuf;
    PrettyWriter<StringBuffer> writer(strbuf);
    it->second.callback()(arguments, &document);
    document.Accept(writer);
    output << strbuf.GetString();
    send_html_headers = false;
  } else {
    it->second.callback()(arguments, &document);
    if (arguments.find("raw") != arguments.end()) {
      document.AddMember(ENABLE_RAW_JSON_KEY, "true", document.GetAllocator());
    }
    if (document.HasMember(ENABLE_RAW_JSON_KEY)) {
      send_html_headers = false;
    }

    const string& full_template_path =
        Substitute("$0/$1/$2", FLAGS_webserver_doc_root, DOC_FOLDER,
            it->second.template_filename());
    ifstream tmpl(full_template_path.c_str());
    if (!tmpl.is_open()) {
      output << "Could not open template: " << full_template_path;
      send_html_headers = false;
    } else {
      stringstream buffer;
      buffer << tmpl.rdbuf();
      RenderTemplate(buffer.str(), Substitute("$0/", FLAGS_webserver_doc_root), document,
          &output);
    }
  }

  VLOG(3) << "Rendering page " << request_info->uri << " took "
          << PrettyPrinter::Print(sw.ElapsedTime(), TCounterType::CPU_TICKS);

  const string& str = output.str();
  // Without styling, render the page as plain text
  if (send_html_headers) {
    sq_printf(connection, "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/html\r\n"
        "Content-Length: %d\r\n"
        "\r\n", (int)str.length());
  } else {
    sq_printf(connection, "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/plain\r\n"
        "Content-Length: %d\r\n"
        "\r\n", (int)str.length());
  }

  // Make sure to use sq_write for printing the body; sq_printf truncates at 8kb
  sq_write(connection, str.c_str(), str.length());
  return PROCESSING_COMPLETE;
}

void Webserver::RegisterUrlCallback(const string& path,
    const string& template_filename, const UrlCallback& callback, bool is_on_nav_bar) {
  upgrade_lock<shared_mutex> lock(url_handlers_lock_);
  upgrade_to_unique_lock<shared_mutex> writer_lock(lock);
  DCHECK(url_handlers_.find(path) == url_handlers_.end())
      << "Duplicate Url handler for: " << path;

  url_handlers_.insert(
      make_pair(path, UrlHandler(callback, template_filename, is_on_nav_bar)));
}

}
<commit_msg>IMPALA-1151: Fix use-after-free in webserver<commit_after>// Copyright 2012 Cloudera Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "util/webserver.h"

#include <boost/algorithm/string.hpp>
#include <boost/bind.hpp>
#include <boost/filesystem.hpp>
#include <boost/foreach.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/mem_fn.hpp>
#include <boost/thread/locks.hpp>
#include <gutil/strings/substitute.h>
#include <map>
#include <fstream>
#include <stdio.h>
#include <signal.h>
#include <string>
#include <mustache/mustache.h>
#include <rapidjson/stringbuffer.h>
#include <rapidjson/prettywriter.h>

#include "common/logging.h"
#include "util/cpu-info.h"
#include "util/disk-info.h"
#include "util/mem-info.h"
#include "util/os-info.h"
#include "util/url-coding.h"
#include "util/debug-util.h"
#include "util/stopwatch.h"
#include "rpc/thrift-util.h"

using namespace std;
using namespace boost;
using namespace boost::filesystem;
using namespace google;
using namespace strings;
using namespace rapidjson;
using namespace mustache;

const char* GetDefaultDocumentRoot();

DEFINE_int32(webserver_port, 25000, "Port to start debug webserver on");
DEFINE_string(webserver_interface, "",
    "Interface to start debug webserver on. If blank, webserver binds to 0.0.0.0");
DEFINE_string(webserver_doc_root, GetDefaultDocumentRoot(),
    "Files under <webserver_doc_root>/www are accessible via the debug webserver. "
    "Defaults to $IMPALA_HOME, or if $IMPALA_HOME is not set, disables the document "
    "root");
DEFINE_bool(enable_webserver_doc_root, true,
    "If true, webserver may serve static files from the webserver_doc_root");

DEFINE_string(webserver_certificate_file, "",
    "The location of the debug webserver's SSL certificate file, in .pem format. If "
    "empty, webserver SSL support is not enabled");
DEFINE_string(webserver_authentication_domain, "",
    "Domain used for debug webserver authentication");
DEFINE_string(webserver_password_file, "",
    "(Optional) Location of .htpasswd file containing user names and hashed passwords for"
    " debug webserver authentication");

static const char* DOC_FOLDER = "/www/";
static const int DOC_FOLDER_LEN = strlen(DOC_FOLDER);

// Easy-to-read constants for Squeasel return codes
static const uint32_t PROCESSING_COMPLETE = 1;
static const uint32_t NOT_PROCESSED = 0;

// Standard keys in the json document sent to templates for rendering. Must be kept in
// sync with the templates themselves.
static const char* COMMON_JSON_KEY = "__common__";

// Returns $IMPALA_HOME if set, otherwise /tmp/impala_www
const char* GetDefaultDocumentRoot() {
  stringstream ss;
  char* impala_home = getenv("IMPALA_HOME");
  if (impala_home == NULL) {
    return ""; // Empty document root means don't serve static files
  } else {
    ss << impala_home;
  }

  // Deliberate memory leak, but this should be called exactly once.
  string* str = new string(ss.str());
  return str->c_str();
}

namespace impala {

const char* Webserver::ENABLE_RAW_JSON_KEY = "__raw__";

Webserver::Webserver() : context_(NULL) {
  http_address_ = MakeNetworkAddress(
      FLAGS_webserver_interface.empty() ? "0.0.0.0" : FLAGS_webserver_interface,
      FLAGS_webserver_port);
}

Webserver::Webserver(const int port) : context_(NULL) {
  http_address_ = MakeNetworkAddress("0.0.0.0", port);
}

Webserver::~Webserver() {
  Stop();
}

void Webserver::RootHandler(const Webserver::ArgumentMap& args, Document* document) {
  Value version(GetVersionString().c_str(), document->GetAllocator());
  document->AddMember("version", version, document->GetAllocator());
  Value cpu_info(CpuInfo::DebugString().c_str(), document->GetAllocator());
  document->AddMember("cpu_info", cpu_info, document->GetAllocator());
  Value mem_info(MemInfo::DebugString().c_str(), document->GetAllocator());
  document->AddMember("mem_info", mem_info, document->GetAllocator());
  Value disk_info(DiskInfo::DebugString().c_str(), document->GetAllocator());
  document->AddMember("disk_info", disk_info, document->GetAllocator());
  Value os_info(OsInfo::DebugString().c_str(), document->GetAllocator());
  document->AddMember("os_info", os_info, document->GetAllocator());
  document->AddMember("pid", getpid(), document->GetAllocator());
}

void Webserver::BuildArgumentMap(const string& args, ArgumentMap* output) {
  vector<string> arg_pairs;
  split(arg_pairs, args, is_any_of("&"));

  BOOST_FOREACH(const string& arg_pair, arg_pairs) {
    vector<string> key_value;
    split(key_value, arg_pair, is_any_of("="));
    if (key_value.empty()) continue;

    string key;
    if (!UrlDecode(key_value[0], &key)) continue;
    string value;
    if (!UrlDecode((key_value.size() >= 2 ? key_value[1] : ""), &value)) continue;
    to_lower(key);
    (*output)[key] = value;
  }
}

bool Webserver::IsSecure() const {
  return !FLAGS_webserver_certificate_file.empty();
}

Status Webserver::Start() {
  LOG(INFO) << "Starting webserver on " << http_address_;

  stringstream listening_spec;
  listening_spec << http_address_;

  if (IsSecure()) {
    LOG(INFO) << "Webserver: Enabling HTTPS support";
    // Squeasel makes sockets with 's' suffixes accept SSL traffic only
    listening_spec << "s";
  }
  string listening_str = listening_spec.str();
  vector<const char*> options;

  if (!FLAGS_webserver_doc_root.empty() && FLAGS_enable_webserver_doc_root) {
    LOG(INFO) << "Document root: " << FLAGS_webserver_doc_root;
    options.push_back("document_root");
    options.push_back(FLAGS_webserver_doc_root.c_str());
  } else {
    LOG(INFO)<< "Document root disabled";
  }

  if (IsSecure()) {
    options.push_back("ssl_certificate");
    options.push_back(FLAGS_webserver_certificate_file.c_str());
  }

  if (!FLAGS_webserver_authentication_domain.empty()) {
    options.push_back("authentication_domain");
    options.push_back(FLAGS_webserver_authentication_domain.c_str());
  }

  if (!FLAGS_webserver_password_file.empty()) {
    // Squeasel doesn't log anything if it can't stat the password file (but will if it
    // can't open it, which it tries to do during a request)
    if (!exists(FLAGS_webserver_password_file)) {
      stringstream ss;
      ss << "Webserver: Password file does not exist: " << FLAGS_webserver_password_file;
      return Status(ss.str());
    }
    LOG(INFO) << "Webserver: Password file is " << FLAGS_webserver_password_file;
    options.push_back("global_auth_file");
    options.push_back(FLAGS_webserver_password_file.c_str());
  }

  options.push_back("listening_ports");
  options.push_back(listening_str.c_str());

  // Options must be a NULL-terminated list
  options.push_back(NULL);

  // squeasel ignores SIGCHLD and we need it to run kinit. This means that since
  // squeasel does not reap its own children CGI programs must be avoided.
  // Save the signal handler so we can restore it after squeasel sets it to be ignored.
  sighandler_t sig_chld = signal(SIGCHLD, SIG_DFL);

  sq_callbacks callbacks;
  memset(&callbacks, 0, sizeof(callbacks));
  callbacks.begin_request = &Webserver::BeginRequestCallbackStatic;
  callbacks.log_message = &Webserver::LogMessageCallbackStatic;

  // To work around not being able to pass member functions as C callbacks, we store a
  // pointer to this server in the per-server state, and register a static method as the
  // default callback. That method unpacks the pointer to this and calls the real
  // callback.
  context_ = sq_start(&callbacks, reinterpret_cast<void*>(this), &options[0]);

  // Restore the child signal handler so wait() works properly.
  signal(SIGCHLD, sig_chld);

  if (context_ == NULL) {
    stringstream error_msg;
    error_msg << "Webserver: Could not start on address " << http_address_;
    return Status(error_msg.str());
  }

  UrlCallback default_callback =
      bind<void>(mem_fn(&Webserver::RootHandler), this, _1, _2);

  RegisterUrlCallback("/", "root.tmpl", default_callback);

  LOG(INFO) << "Webserver started";
  return Status::OK;
}

void Webserver::Stop() {
  if (context_ != NULL) {
    sq_stop(context_);
    context_ = NULL;
  }
}

void Webserver::GetCommonJson(Document* document) {
  DCHECK(document != NULL);
  Value obj(kObjectType);
  obj.AddMember("process-name", google::ProgramInvocationShortName(),
      document->GetAllocator());

  Value lst(kArrayType);
  BOOST_FOREACH(const UrlHandlerMap::value_type& handler, url_handlers_) {
    if (handler.second.is_on_nav_bar()) {
      Value obj(kObjectType);
      obj.AddMember("link", handler.first.c_str(), document->GetAllocator());
      obj.AddMember("title", handler.first.c_str(), document->GetAllocator());
      lst.PushBack(obj, document->GetAllocator());
    }
  }

  obj.AddMember("navbar", lst, document->GetAllocator());
  document->AddMember(COMMON_JSON_KEY, obj, document->GetAllocator());
}

int Webserver::LogMessageCallbackStatic(const struct sq_connection* connection,
    const char* message) {
  if (message != NULL) {
    LOG(INFO) << "Webserver: " << message;
  }
  return PROCESSING_COMPLETE;
}

int Webserver::BeginRequestCallbackStatic(struct sq_connection* connection) {
  struct sq_request_info* request_info = sq_get_request_info(connection);
  Webserver* instance = reinterpret_cast<Webserver*>(request_info->user_data);
  return instance->BeginRequestCallback(connection, request_info);
}

int Webserver::BeginRequestCallback(struct sq_connection* connection,
    struct sq_request_info* request_info) {
  if (!FLAGS_webserver_doc_root.empty() && FLAGS_enable_webserver_doc_root) {
    if (strncmp(DOC_FOLDER, request_info->uri, DOC_FOLDER_LEN) == 0) {
      VLOG(2) << "HTTP File access: " << request_info->uri;
      // Let Squeasel deal with this request; returning NULL will fall through
      // to the default handler which will serve files.
      return NOT_PROCESSED;
    }
  }
  shared_lock<shared_mutex> lock(url_handlers_lock_);
  UrlHandlerMap::const_iterator it = url_handlers_.find(request_info->uri);
  if (it == url_handlers_.end()) {
    sq_printf(connection, "HTTP/1.1 404 Not Found\r\n"
        "Content-Type: text/plain\r\n\r\n");
    sq_printf(connection, "No handler for URI %s\r\n\r\n", request_info->uri);
    return PROCESSING_COMPLETE;
  }

  map<string, string> arguments;
  if (request_info->query_string != NULL) {
    BuildArgumentMap(request_info->query_string, &arguments);
  }
  MonotonicStopWatch sw;
  sw.Start();

  Document document;
  document.SetObject();
  GetCommonJson(&document);

  bool send_html_headers = true;
  // The output of this page is accumulated into this stringstream.
  stringstream output;
  bool raw_json = (arguments.find("json") != arguments.end());
  if (raw_json) {
    // Callbacks may optionally be rendered as a text-only, pretty-printed Json document
    // (mostly for debugging or integration with third-party tools).
    StringBuffer strbuf;
    PrettyWriter<StringBuffer> writer(strbuf);
    it->second.callback()(arguments, &document);
    document.Accept(writer);
    output << strbuf.GetString();
    send_html_headers = false;
  } else {
    it->second.callback()(arguments, &document);
    if (arguments.find("raw") != arguments.end()) {
      document.AddMember(ENABLE_RAW_JSON_KEY, "true", document.GetAllocator());
    }
    if (document.HasMember(ENABLE_RAW_JSON_KEY)) {
      send_html_headers = false;
    }

    const string& full_template_path =
        Substitute("$0/$1/$2", FLAGS_webserver_doc_root, DOC_FOLDER,
            it->second.template_filename());
    ifstream tmpl(full_template_path.c_str());
    if (!tmpl.is_open()) {
      output << "Could not open template: " << full_template_path;
      send_html_headers = false;
    } else {
      stringstream buffer;
      buffer << tmpl.rdbuf();
      RenderTemplate(buffer.str(), Substitute("$0/", FLAGS_webserver_doc_root), document,
          &output);
    }
  }

  VLOG(3) << "Rendering page " << request_info->uri << " took "
          << PrettyPrinter::Print(sw.ElapsedTime(), TCounterType::CPU_TICKS);

  const string& str = output.str();
  // Without styling, render the page as plain text
  if (send_html_headers) {
    sq_printf(connection, "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/html\r\n"
        "Content-Length: %d\r\n"
        "\r\n", (int)str.length());
  } else {
    sq_printf(connection, "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/plain\r\n"
        "Content-Length: %d\r\n"
        "\r\n", (int)str.length());
  }

  // Make sure to use sq_write for printing the body; sq_printf truncates at 8kb
  sq_write(connection, str.c_str(), str.length());
  return PROCESSING_COMPLETE;
}

void Webserver::RegisterUrlCallback(const string& path,
    const string& template_filename, const UrlCallback& callback, bool is_on_nav_bar) {
  upgrade_lock<shared_mutex> lock(url_handlers_lock_);
  upgrade_to_unique_lock<shared_mutex> writer_lock(lock);
  DCHECK(url_handlers_.find(path) == url_handlers_.end())
      << "Duplicate Url handler for: " << path;

  url_handlers_.insert(
      make_pair(path, UrlHandler(callback, template_filename, is_on_nav_bar)));
}

}
<|endoftext|>